[Mlir-commits] [mlir] [MLIR] Split ExecutionEngine Initialization out of ctor into an explicit method call (PR #153524)
Shenghang Tsai
llvmlistbot at llvm.org
Sat Aug 16 22:22:37 PDT 2025
https://github.com/jackalcooper updated https://github.com/llvm/llvm-project/pull/153524
>From 975a39e9d5f9a89a3bcbcdea87221500c9365eec Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Thu, 14 Aug 2025 08:22:51 +0800
Subject: [PATCH 01/15] port changes from last PR
---
mlir/include/mlir-c/ExecutionEngine.h | 7 +++
.../mlir/ExecutionEngine/ExecutionEngine.h | 9 +++
.../Bindings/Python/ExecutionEngineModule.cpp | 13 ++++-
.../CAPI/ExecutionEngine/ExecutionEngine.cpp | 9 ++-
mlir/lib/ExecutionEngine/ExecutionEngine.cpp | 20 ++++---
mlir/lib/ExecutionEngine/JitRunner.cpp | 2 +
mlir/test/CAPI/execution_engine.c | 55 +++++++++++++++++++
mlir/unittests/ExecutionEngine/Invoke.cpp | 38 +++++++++++++
8 files changed, 141 insertions(+), 12 deletions(-)
diff --git a/mlir/include/mlir-c/ExecutionEngine.h b/mlir/include/mlir-c/ExecutionEngine.h
index 99cddc5c2598d..1a58d68533f24 100644
--- a/mlir/include/mlir-c/ExecutionEngine.h
+++ b/mlir/include/mlir-c/ExecutionEngine.h
@@ -46,6 +46,13 @@ MLIR_CAPI_EXPORTED MlirExecutionEngine mlirExecutionEngineCreate(
MlirModule op, int optLevel, int numPaths,
const MlirStringRef *sharedLibPaths, bool enableObjectDump);
+/// 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.
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..5bd71d68d253a 100644
--- a/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
+++ b/mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
@@ -227,6 +227,13 @@ class ExecutionEngine {
llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
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 resolvable before initialization by calling
+ /// `registerSymbols` or including shared libraries.
+ void initialize();
+
private:
/// Ordering of llvmContext and jit is important for destruction purposes: the
/// jit must be destroyed before the context.
@@ -250,6 +257,8 @@ class ExecutionEngine {
/// Destroy functions in the libraries loaded by the ExecutionEngine that are
/// called when this ExecutionEngine is destructed.
SmallVector<LibraryDestroyFn> destroyFns;
+
+ bool isInitialized = false;
};
} // namespace mlir
diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index 81dada3553622..4f7a4a628e246 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;
@@ -124,6 +124,17 @@ 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. 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) {
diff --git a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
index 306cebd236be9..2dbb993b1640f 100644
--- a/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/CAPI/ExecutionEngine/ExecutionEngine.cpp
@@ -68,6 +68,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 +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 f704fbfbe8fff..52162a43aeae3 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,6 @@ 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()));
-
return std::move(engine);
}
@@ -442,6 +435,7 @@ Expected<void *> ExecutionEngine::lookup(StringRef name) const {
Error ExecutionEngine::invokePacked(StringRef name,
MutableArrayRef<void *> args) {
+ initialize();
auto expectedFPtr = lookupPacked(name);
if (!expectedFPtr)
return expectedFPtr.takeError();
@@ -451,3 +445,13 @@ Error ExecutionEngine::invokePacked(StringRef name,
return Error::success();
}
+
+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()));
+ isInitialized = true;
+}
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 4751288c3ee4b..900588aeea2da 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);
+
+ if (mlirExecutionEngineIsNull(jit)) {
+ fprintf(stderr, "Execution engine creation failed");
+ exit(2);
+ }
+
+ // Register callback symbol before initialization
+ mlirExecutionEngineRegisterSymbol(
+ jit, mlirStringRefCreateFromCString("_mlir_ciface_init_callback"),
+ (void *)(uintptr_t)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;
}
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index 312b10f28143f..712d4c49a2afe 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -322,4 +322,42 @@ 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<ModuleOp>(moduleStr, &context);
+ ASSERT_TRUE(!!module);
+ ASSERT_TRUE(succeeded(lowerToLLVMDialect(*module)));
+ ExecutionEngineOptions jitOptions;
+ 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 793e83a840d5055fb87cc0c7dd12a393e48b6628 Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Thu, 14 Aug 2025 08:42:40 +0800
Subject: [PATCH 02/15] update GPU docs to reflect latest impl
---
mlir/docs/Dialects/GPU.md | 10 +++++++++-
1 file changed, 9 insertions(+), 1 deletion(-)
diff --git a/mlir/docs/Dialects/GPU.md b/mlir/docs/Dialects/GPU.md
index 94b053daa1615..3578f430ec486 100644
--- a/mlir/docs/Dialects/GPU.md
+++ b/mlir/docs/Dialects/GPU.md
@@ -193,10 +193,18 @@ llvm.func @foo() {
// mlir-translate --mlir-to-llvmir:
@binary_bin_cst = internal constant [6 x i8] c"AMDGPU", align 8
@binary_func_kernel_name = private unnamed_addr constant [7 x i8] c"func\00", align 1
+ at module = internal global ptr null
+ at llvm.global_ctors = appending global [1 x {i32, ptr, ptr}] [{i32 123, ptr @binary_load, ptr null}]
+ at llvm.global_dtors = appending global [1 x {i32, ptr, ptr}] [{i32 123, ptr @binary_unload, ptr null}]
+define internal void @binary_load() section ".text.startup" {
+entry:
+ %0 = call ptr @mgpuModuleLoad(ptr @binary_bin_cst)
+ store ptr %0, ptr @module
+ ...
+}
...
define void @foo() {
...
- %module = call ptr @mgpuModuleLoad(ptr @binary_bin_cst)
%kernel = call ptr @mgpuModuleGetFunction(ptr %module, ptr @binary_func_kernel_name)
call void @mgpuLaunchKernel(ptr %kernel, ...) ; Launch the kernel
...
>From e948795c900104a28e8cb51c98a45a8fbe69f622 Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Thu, 14 Aug 2025 08:47:08 +0800
Subject: [PATCH 03/15] add module symbol->ptr load
---
mlir/docs/Dialects/GPU.md | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/docs/Dialects/GPU.md b/mlir/docs/Dialects/GPU.md
index 3578f430ec486..a06183ad39dce 100644
--- a/mlir/docs/Dialects/GPU.md
+++ b/mlir/docs/Dialects/GPU.md
@@ -193,7 +193,7 @@ llvm.func @foo() {
// mlir-translate --mlir-to-llvmir:
@binary_bin_cst = internal constant [6 x i8] c"AMDGPU", align 8
@binary_func_kernel_name = private unnamed_addr constant [7 x i8] c"func\00", align 1
- at module = internal global ptr null
+ at binary_@module = internal global ptr null
@llvm.global_ctors = appending global [1 x {i32, ptr, ptr}] [{i32 123, ptr @binary_load, ptr null}]
@llvm.global_dtors = appending global [1 x {i32, ptr, ptr}] [{i32 123, ptr @binary_unload, ptr null}]
define internal void @binary_load() section ".text.startup" {
@@ -205,6 +205,7 @@ entry:
...
define void @foo() {
...
+ %module = load ptr, ptr @binary_module, align 8
%kernel = call ptr @mgpuModuleGetFunction(ptr %module, ptr @binary_func_kernel_name)
call void @mgpuLaunchKernel(ptr %kernel, ...) ; Launch the kernel
...
>From fe4d555d7b6bdc8c379416aa3ca2ab26b89b568c Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Thu, 14 Aug 2025 09:11:10 +0800
Subject: [PATCH 04/15] __has_feature(address_sanitizer)
---
mlir/unittests/ExecutionEngine/Invoke.cpp | 5 +++++
1 file changed, 5 insertions(+)
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index 712d4c49a2afe..566c587957992 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -326,6 +326,11 @@ static int initCnt = 0;
// A helper function that will be called during the JIT's initialization.
static void initCallback() { initCnt += 1; }
+#if __has_feature(memory_sanitizer) || __has_feature(address_sanitizer)
+#define MAYBE_JITCallback DISABLED_JITCallback
+#else
+#define MAYBE_JITCallback SKIP_WITHOUT_JIT(JITCallback)
+#endif
TEST(GlobalCtorJit, MAYBE_JITCallback) {
std::string moduleStr = R"mlir(
llvm.mlir.global_ctors ctors = [@ctor], priorities = [0 : i32], data = [#llvm.zero]
>From 0a55a9b4070e2eaa039272587daae1ed0390265b Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Thu, 14 Aug 2025 09:25:55 +0800
Subject: [PATCH 05/15] remove test in C
---
mlir/test/CAPI/execution_engine.c | 55 -------------------------------
1 file changed, 55 deletions(-)
diff --git a/mlir/test/CAPI/execution_engine.c b/mlir/test/CAPI/execution_engine.c
index 900588aeea2da..4751288c3ee4b 100644
--- a/mlir/test/CAPI/execution_engine.c
+++ b/mlir/test/CAPI/execution_engine.c
@@ -137,60 +137,6 @@ 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);
-
- if (mlirExecutionEngineIsNull(jit)) {
- fprintf(stderr, "Execution engine creation failed");
- exit(2);
- }
-
- // Register callback symbol before initialization
- mlirExecutionEngineRegisterSymbol(
- jit, mlirStringRefCreateFromCString("_mlir_ciface_init_callback"),
- (void *)(uintptr_t)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
@@ -201,6 +147,5 @@ int main(void) {
TEST(testSimpleExecution);
TEST(testOmpCreation);
- TEST(testGlobalCtorJitCallback);
return 0;
}
>From 637e13faeb9ca32d1072ccd8acace35d398296fe Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Thu, 14 Aug 2025 09:32:29 +0800
Subject: [PATCH 06/15] better test naming
---
mlir/unittests/ExecutionEngine/Invoke.cpp | 6 +++---
1 file changed, 3 insertions(+), 3 deletions(-)
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index 566c587957992..63ff6f44ed8a9 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -327,11 +327,11 @@ static int initCnt = 0;
static void initCallback() { initCnt += 1; }
#if __has_feature(memory_sanitizer) || __has_feature(address_sanitizer)
-#define MAYBE_JITCallback DISABLED_JITCallback
+#define MAYBE_JITCallbackInGlobalCtor DISABLED_JITCallbackInGlobalCtor
#else
-#define MAYBE_JITCallback SKIP_WITHOUT_JIT(JITCallback)
+#define MAYBE_JITCallbackInGlobalCtor SKIP_WITHOUT_JIT(JITCallbackInGlobalCtor)
#endif
-TEST(GlobalCtorJit, MAYBE_JITCallback) {
+TEST(MLIRExecutionEngine, MAYBE_JITCallbackInGlobalCtor) {
std::string moduleStr = R"mlir(
llvm.mlir.global_ctors ctors = [@ctor], priorities = [0 : i32], data = [#llvm.zero]
llvm.func @ctor() {
>From 25db89b31c4d5abb924d4db4c1d8e3b8e1a7a858 Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Thu, 14 Aug 2025 10:09:45 +0800
Subject: [PATCH 07/15] skip asserts for AArch64
---
mlir/unittests/ExecutionEngine/Invoke.cpp | 14 +++++++++++++-
1 file changed, 13 insertions(+), 1 deletion(-)
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index 63ff6f44ed8a9..8d5f99ec71c13 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -352,6 +352,8 @@ TEST(MLIRExecutionEngine, MAYBE_JITCallbackInGlobalCtor) {
ExecutionEngineOptions jitOptions;
auto jitOrError = ExecutionEngine::create(*module, jitOptions);
ASSERT_TRUE(!!jitOrError);
+ // validate initialization is not run on construction
+ ASSERT_EQ(initCnt, 0);
auto jit = std::move(jitOrError.get());
// Define any extra symbols so they're available at initialization.
jit->registerSymbols([&](llvm::orc::MangleAndInterner interner) {
@@ -362,7 +364,17 @@ TEST(MLIRExecutionEngine, MAYBE_JITCallbackInGlobalCtor) {
return symbolMap;
});
jit->initialize();
- ASSERT_EQ(initCnt, 1);
+ // TODO: Allow JIT initialize for AArch64. Currently there's a bug causing a
+ // crash for AArch64 see related issue #71963.
+ auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
+ ASSERT_TRUE(!!tmBuilderOrError);
+ if (!tmBuilderOrError->getTargetTriple().isAArch64()) {
+ // validate the side effect of initialization
+ ASSERT_EQ(initCnt, 1);
+ // next initialization should be noop
+ jit->initialize();
+ ASSERT_EQ(initCnt, 1);
+ }
}
#endif // _WIN32
>From 413551072e782e9f3157482f381fc37f359054ca Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Thu, 14 Aug 2025 10:25:44 +0800
Subject: [PATCH 08/15] add more sanitizers to skip
---
mlir/unittests/ExecutionEngine/Invoke.cpp | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index 8d5f99ec71c13..f827ec9b77066 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -326,7 +326,9 @@ static int initCnt = 0;
// A helper function that will be called during the JIT's initialization.
static void initCallback() { initCnt += 1; }
-#if __has_feature(memory_sanitizer) || __has_feature(address_sanitizer)
+#if __has_feature(memory_sanitizer) || __has_feature(address_sanitizer) || \
+ __has_feature(hwaddress_sanitizer) || \
+ __has_feature(undefined_behavior_sanitizer)
#define MAYBE_JITCallbackInGlobalCtor DISABLED_JITCallbackInGlobalCtor
#else
#define MAYBE_JITCallbackInGlobalCtor SKIP_WITHOUT_JIT(JITCallbackInGlobalCtor)
>From 03a62ac8b779726dbc397ea81ccc848957c54710 Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Fri, 15 Aug 2025 17:19:23 +0800
Subject: [PATCH 09/15] fix typo
---
mlir/docs/Dialects/GPU.md | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/docs/Dialects/GPU.md b/mlir/docs/Dialects/GPU.md
index a06183ad39dce..c74c8cb3f036e 100644
--- a/mlir/docs/Dialects/GPU.md
+++ b/mlir/docs/Dialects/GPU.md
@@ -193,7 +193,7 @@ llvm.func @foo() {
// mlir-translate --mlir-to-llvmir:
@binary_bin_cst = internal constant [6 x i8] c"AMDGPU", align 8
@binary_func_kernel_name = private unnamed_addr constant [7 x i8] c"func\00", align 1
- at binary_@module = internal global ptr null
+ at binary_module = internal global ptr null
@llvm.global_ctors = appending global [1 x {i32, ptr, ptr}] [{i32 123, ptr @binary_load, ptr null}]
@llvm.global_dtors = appending global [1 x {i32, ptr, ptr}] [{i32 123, ptr @binary_unload, ptr null}]
define internal void @binary_load() section ".text.startup" {
>From b5a4244d598ee6e1ac4ad6b45a410f156d35a804 Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Fri, 15 Aug 2025 17:43:15 +0800
Subject: [PATCH 10/15] py version test
---
.../Bindings/Python/ExecutionEngineModule.cpp | 2 +-
.../mlir/_mlir_libs/_mlirExecutionEngine.pyi | 1 +
mlir/test/python/execution_engine.py | 37 +++++++++++++++++++
3 files changed, 39 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
index 4f7a4a628e246..4885d62c56e6e 100644
--- a/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
+++ b/mlir/lib/Bindings/Python/ExecutionEngineModule.cpp
@@ -133,7 +133,7 @@ NB_MODULE(_mlirExecutionEngine, m) {
"`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 "
+ "initialization by calling `register_runtime` or including "
"shared libraries.")
.def(
"dump_to_object_file",
diff --git a/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi b/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi
index 58d453d2b2d37..4b82c78489295 100644
--- a/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi
+++ b/mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi
@@ -19,5 +19,6 @@ class ExecutionEngine:
def dump_to_object_file(self, file_name: str) -> None: ...
def raw_lookup(self, func_name: str) -> int: ...
def raw_register_runtime(self, name: str, callback: object) -> None: ...
+ def init() -> None: ...
@property
def _CAPIPtr(self) -> object: ...
diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index d569fcef32bfd..6e8a422b12d29 100644
--- a/mlir/test/python/execution_engine.py
+++ b/mlir/test/python/execution_engine.py
@@ -339,6 +339,43 @@ def callback(a):
run(testUnrankedMemRefWithOffsetCallback)
+# Test JIT callback in global constructor
+# CHECK-LABEL: TEST: testJITCallbackInGlobalCtor
+def testJITCallbackInGlobalCtor():
+ init_cnt = 0
+
+ @ctypes.CFUNCTYPE(None)
+ def initCallback():
+ nonlocal init_cnt
+ init_cnt += 1
+
+ with Context():
+ module = Module.parse(r"""
+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 }
+ """)
+
+ # Setup execution engine
+ execution_engine = ExecutionEngine(lowerToLLVM(module))
+
+ # Validate initialization hasn't run yet
+ assert init_cnt == 0
+
+ # # Register callback
+ execution_engine.register_runtime("init_callback", initCallback)
+
+ # # Initialize and verify
+ execution_engine.initialize()
+ assert init_cnt == 1
+ # # Second initialization should be no-op
+ execution_engine.initialize()
+ assert init_cnt == 1
+
+run(testJITCallbackInGlobalCtor)
# Test addition of two memrefs.
# CHECK-LABEL: TEST: testMemrefAdd
>From acf71c4d3f0caccddf371bbc97f225c4466ef34c Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Fri, 15 Aug 2025 17:48:31 +0800
Subject: [PATCH 11/15] doc @binary_unload()
---
mlir/docs/Dialects/GPU.md | 8 +++++++-
1 file changed, 7 insertions(+), 1 deletion(-)
diff --git a/mlir/docs/Dialects/GPU.md b/mlir/docs/Dialects/GPU.md
index c74c8cb3f036e..8d4d2ca3e5743 100644
--- a/mlir/docs/Dialects/GPU.md
+++ b/mlir/docs/Dialects/GPU.md
@@ -199,7 +199,13 @@ llvm.func @foo() {
define internal void @binary_load() section ".text.startup" {
entry:
%0 = call ptr @mgpuModuleLoad(ptr @binary_bin_cst)
- store ptr %0, ptr @module
+ store ptr %0, ptr @binary_module
+ ...
+}
+define internal void @binary_unload() section ".text.startup" {
+entry:
+ %0 = load ptr, ptr @binary_module, align 8
+ call void @mgpuModuleUnload(ptr %0)
...
}
...
>From 5dd9b203dfac58d1c388673c910e5b8defa869af Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Fri, 15 Aug 2025 18:02:34 +0800
Subject: [PATCH 12/15] fmt
---
mlir/test/python/execution_engine.py | 23 ++++++++++++-----------
1 file changed, 12 insertions(+), 11 deletions(-)
diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index 6e8a422b12d29..e67863f5b3694 100644
--- a/mlir/test/python/execution_engine.py
+++ b/mlir/test/python/execution_engine.py
@@ -382,17 +382,18 @@ def initCallback():
def testMemrefAdd():
with Context():
module = Module.parse(
- """
- module {
- func.func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {
- %0 = arith.constant 0 : index
- %1 = memref.load %arg0[%0] : memref<1xf32>
- %2 = memref.load %arg1[] : memref<f32>
- %3 = arith.addf %1, %2 : f32
- memref.store %3, %arg2[%0] : memref<1xf32>
- return
- }
- } """
+ r"""
+module {
+ func.func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {
+ %0 = arith.constant 0 : index
+ %1 = memref.load %arg0[%0] : memref<1xf32>
+ %2 = memref.load %arg1[] : memref<f32>
+ %3 = arith.addf %1, %2 : f32
+ memref.store %3, %arg2[%0] : memref<1xf32>
+ return
+ }
+}
+ """
)
arg1 = np.array([32.5]).astype(np.float32)
arg2 = np.array(6).astype(np.float32)
>From 6a78eeafdd7d6de258d08d292eb6455ebdbb9631 Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Sun, 17 Aug 2025 12:56:18 +0800
Subject: [PATCH 13/15] restore c tests
---
mlir/test/CAPI/CMakeLists.txt | 7 ++
mlir/test/CAPI/global_constructors.c | 113 +++++++++++++++++++++++++++
mlir/test/CMakeLists.txt | 1 +
mlir/test/lit.cfg.py | 1 +
4 files changed, 122 insertions(+)
create mode 100644 mlir/test/CAPI/global_constructors.c
diff --git a/mlir/test/CAPI/CMakeLists.txt b/mlir/test/CAPI/CMakeLists.txt
index a7f9eb9b4efe8..d45142510a496 100644
--- a/mlir/test/CAPI/CMakeLists.txt
+++ b/mlir/test/CAPI/CMakeLists.txt
@@ -30,6 +30,13 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
MLIRCAPIConversion
MLIRCAPIExecutionEngine
MLIRCAPIRegisterEverything
+)
+ _add_capi_test_executable(mlir-capi-global-constructors-test
+ global_constructors.c
+ LINK_LIBS PRIVATE
+ MLIRCAPIConversion
+ MLIRCAPIExecutionEngine
+ MLIRCAPIRegisterEverything
)
endif()
diff --git a/mlir/test/CAPI/global_constructors.c b/mlir/test/CAPI/global_constructors.c
new file mode 100644
index 0000000000000..bd2fe1416f0df
--- /dev/null
+++ b/mlir/test/CAPI/global_constructors.c
@@ -0,0 +1,113 @@
+//===- global_constructors.c - Test JIT with the global constructors ------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM
+// Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+// UNSUPPORTED: target=aarch64{{.*}}, target=arm64{{.*}}
+/* RUN: mlir-capi-global-constructors-test 2>&1 | FileCheck %s
+ */
+/* REQUIRES: host-supports-jit
+ */
+
+#include "mlir-c/Conversion.h"
+#include "mlir-c/ExecutionEngine.h"
+#include "mlir-c/IR.h"
+#include "mlir-c/RegisterEverything.h"
+
+#include <assert.h>
+#include <math.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+static void registerAllUpstreamDialects(MlirContext ctx) {
+ MlirDialectRegistry registry = mlirDialectRegistryCreate();
+ mlirRegisterAllDialects(registry);
+ mlirContextAppendDialectRegistry(ctx, registry);
+ mlirDialectRegistryDestroy(registry);
+}
+
+void lowerModuleToLLVM(MlirContext ctx, MlirModule module) {
+ MlirPassManager pm = mlirPassManagerCreate(ctx);
+ MlirOpPassManager opm = mlirPassManagerGetNestedUnder(
+ pm, mlirStringRefCreateFromCString("func.func"));
+ mlirPassManagerAddOwnedPass(pm, mlirCreateConversionConvertFuncToLLVMPass());
+ mlirOpPassManagerAddOwnedPass(
+ opm, mlirCreateConversionArithToLLVMConversionPass());
+ MlirLogicalResult status =
+ mlirPassManagerRunOnOp(pm, mlirModuleGetOperation(module));
+ if (mlirLogicalResultIsFailure(status)) {
+ fprintf(stderr, "Unexpected failure running pass pipeline\n");
+ exit(2);
+ }
+ mlirPassManagerDestroy(pm);
+}
+
+// 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);
+
+ if (mlirExecutionEngineIsNull(jit)) {
+ fprintf(stderr, "Execution engine creation failed");
+ exit(2);
+ }
+
+ // Register callback symbol before initialization
+ mlirExecutionEngineRegisterSymbol(
+ jit, mlirStringRefCreateFromCString("_mlir_ciface_init_callback"),
+ (void *)(uintptr_t)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
+#define STRINGIFY(x) _STRINGIFY(x)
+#define TEST(test) \
+ printf("Running test '" STRINGIFY(test) "'\n"); \
+ test();
+ TEST(testGlobalCtorJitCallback);
+ return 0;
+}
diff --git a/mlir/test/CMakeLists.txt b/mlir/test/CMakeLists.txt
index a4a942de3c9a7..0b98eaaf3391c 100644
--- a/mlir/test/CMakeLists.txt
+++ b/mlir/test/CMakeLists.txt
@@ -141,6 +141,7 @@ if(LLVM_ENABLE_PIC AND TARGET ${LLVM_NATIVE_ARCH})
llc
mlir_async_runtime
mlir-capi-execution-engine-test
+ mlir-capi-global-constructors-test
mlir_c_runner_utils
mlir_runner_utils
mlir_float16_utils
diff --git a/mlir/test/lit.cfg.py b/mlir/test/lit.cfg.py
index ba7eeeed8ef3f..5e9347d784b38 100644
--- a/mlir/test/lit.cfg.py
+++ b/mlir/test/lit.cfg.py
@@ -190,6 +190,7 @@ def find_real_python_interpreter():
"mlir-translate",
"mlir-lsp-server",
"mlir-capi-execution-engine-test",
+ "mlir-capi-global-constructors-test",
"mlir-capi-ir-test",
"mlir-capi-irdl-test",
"mlir-capi-llvm-test",
>From 1a0e331f4ef95f0c4f5728070bdd8cee36268672 Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Sun, 17 Aug 2025 13:07:33 +0800
Subject: [PATCH 14/15] refine C++ and Python tests
---
mlir/test/python/execution_engine.py | 60 ++++---------------
mlir/test/python/global_constructors.py | 70 +++++++++++++++++++++++
mlir/unittests/ExecutionEngine/Invoke.cpp | 30 ++++------
3 files changed, 93 insertions(+), 67 deletions(-)
create mode 100644 mlir/test/python/global_constructors.py
diff --git a/mlir/test/python/execution_engine.py b/mlir/test/python/execution_engine.py
index e67863f5b3694..d569fcef32bfd 100644
--- a/mlir/test/python/execution_engine.py
+++ b/mlir/test/python/execution_engine.py
@@ -339,61 +339,23 @@ def callback(a):
run(testUnrankedMemRefWithOffsetCallback)
-# Test JIT callback in global constructor
-# CHECK-LABEL: TEST: testJITCallbackInGlobalCtor
-def testJITCallbackInGlobalCtor():
- init_cnt = 0
-
- @ctypes.CFUNCTYPE(None)
- def initCallback():
- nonlocal init_cnt
- init_cnt += 1
-
- with Context():
- module = Module.parse(r"""
-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 }
- """)
-
- # Setup execution engine
- execution_engine = ExecutionEngine(lowerToLLVM(module))
-
- # Validate initialization hasn't run yet
- assert init_cnt == 0
-
- # # Register callback
- execution_engine.register_runtime("init_callback", initCallback)
-
- # # Initialize and verify
- execution_engine.initialize()
- assert init_cnt == 1
- # # Second initialization should be no-op
- execution_engine.initialize()
- assert init_cnt == 1
-
-run(testJITCallbackInGlobalCtor)
# Test addition of two memrefs.
# CHECK-LABEL: TEST: testMemrefAdd
def testMemrefAdd():
with Context():
module = Module.parse(
- r"""
-module {
- func.func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {
- %0 = arith.constant 0 : index
- %1 = memref.load %arg0[%0] : memref<1xf32>
- %2 = memref.load %arg1[] : memref<f32>
- %3 = arith.addf %1, %2 : f32
- memref.store %3, %arg2[%0] : memref<1xf32>
- return
- }
-}
- """
+ """
+ module {
+ func.func @main(%arg0: memref<1xf32>, %arg1: memref<f32>, %arg2: memref<1xf32>) attributes { llvm.emit_c_interface } {
+ %0 = arith.constant 0 : index
+ %1 = memref.load %arg0[%0] : memref<1xf32>
+ %2 = memref.load %arg1[] : memref<f32>
+ %3 = arith.addf %1, %2 : f32
+ memref.store %3, %arg2[%0] : memref<1xf32>
+ return
+ }
+ } """
)
arg1 = np.array([32.5]).astype(np.float32)
arg2 = np.array(6).astype(np.float32)
diff --git a/mlir/test/python/global_constructors.py b/mlir/test/python/global_constructors.py
new file mode 100644
index 0000000000000..074ecd867ac7d
--- /dev/null
+++ b/mlir/test/python/global_constructors.py
@@ -0,0 +1,70 @@
+# UNSUPPORTED: target=aarch64{{.*}}, target=arm64{{.*}}
+# RUN: %PYTHON %s 2>&1 | FileCheck %s
+# REQUIRES: host-supports-jit
+import gc, sys, os, tempfile
+from mlir.ir import *
+from mlir.passmanager import *
+from mlir.execution_engine import *
+from mlir.runtime import *
+
+
+# Log everything to stderr and flush so that we have a unified stream to match
+# errors/info emitted by MLIR to stderr.
+def log(*args):
+ print(*args, file=sys.stderr)
+ sys.stderr.flush()
+
+
+def run(f):
+ log("\nTEST:", f.__name__)
+ f()
+ gc.collect()
+ assert Context._get_live_count() == 0
+
+def lowerToLLVM(module):
+ pm = PassManager.parse(
+ "builtin.module(convert-func-to-llvm,reconcile-unrealized-casts)"
+ )
+ pm.run(module.operation)
+ return module
+
+# Test JIT callback in global constructor
+# CHECK-LABEL: TEST: testJITCallbackInGlobalCtor
+def testJITCallbackInGlobalCtor():
+ init_cnt = 0
+
+ @ctypes.CFUNCTYPE(None)
+ def initCallback():
+ nonlocal init_cnt
+ init_cnt += 1
+
+ with Context():
+ module = Module.parse(
+ r"""
+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 }
+ """
+ )
+
+ # Setup execution engine
+ execution_engine = ExecutionEngine(lowerToLLVM(module))
+
+ # Validate initialization hasn't run yet
+ assert init_cnt == 0
+
+ # # Register callback
+ execution_engine.register_runtime("init_callback", initCallback)
+
+ # # Initialize and verify
+ execution_engine.initialize()
+ assert init_cnt == 1
+ # # Second initialization should be no-op
+ execution_engine.initialize()
+ assert init_cnt == 1
+
+
+run(testJITCallbackInGlobalCtor)
diff --git a/mlir/unittests/ExecutionEngine/Invoke.cpp b/mlir/unittests/ExecutionEngine/Invoke.cpp
index f827ec9b77066..7cf3e26ffeab5 100644
--- a/mlir/unittests/ExecutionEngine/Invoke.cpp
+++ b/mlir/unittests/ExecutionEngine/Invoke.cpp
@@ -326,14 +326,14 @@ static int initCnt = 0;
// A helper function that will be called during the JIT's initialization.
static void initCallback() { initCnt += 1; }
-#if __has_feature(memory_sanitizer) || __has_feature(address_sanitizer) || \
- __has_feature(hwaddress_sanitizer) || \
- __has_feature(undefined_behavior_sanitizer)
-#define MAYBE_JITCallbackInGlobalCtor DISABLED_JITCallbackInGlobalCtor
-#else
-#define MAYBE_JITCallbackInGlobalCtor SKIP_WITHOUT_JIT(JITCallbackInGlobalCtor)
-#endif
TEST(MLIRExecutionEngine, MAYBE_JITCallbackInGlobalCtor) {
+ // TODO: Allow JIT initialize for AArch64. Currently there's a bug causing a
+ // crash for AArch64 see related issue #71963.
+ auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
+ ASSERT_TRUE(!!tmBuilderOrError);
+ if (!tmBuilderOrError->getTargetTriple().isAArch64()) {
+ return;
+ }
std::string moduleStr = R"mlir(
llvm.mlir.global_ctors ctors = [@ctor], priorities = [0 : i32], data = [#llvm.zero]
llvm.func @ctor() {
@@ -366,17 +366,11 @@ TEST(MLIRExecutionEngine, MAYBE_JITCallbackInGlobalCtor) {
return symbolMap;
});
jit->initialize();
- // TODO: Allow JIT initialize for AArch64. Currently there's a bug causing a
- // crash for AArch64 see related issue #71963.
- auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
- ASSERT_TRUE(!!tmBuilderOrError);
- if (!tmBuilderOrError->getTargetTriple().isAArch64()) {
- // validate the side effect of initialization
- ASSERT_EQ(initCnt, 1);
- // next initialization should be noop
- jit->initialize();
- ASSERT_EQ(initCnt, 1);
- }
+ // validate the side effect of initialization
+ ASSERT_EQ(initCnt, 1);
+ // next initialization should be noop
+ jit->initialize();
+ ASSERT_EQ(initCnt, 1);
}
#endif // _WIN32
>From c8ac9dc9f77f6b4a7fc3e478746fadd5d8305015 Mon Sep 17 00:00:00 2001
From: jackalcooper <jackalcooper at gmail.com>
Date: Sun, 17 Aug 2025 13:22:15 +0800
Subject: [PATCH 15/15] fix fmt
---
mlir/test/python/global_constructors.py | 2 ++
1 file changed, 2 insertions(+)
diff --git a/mlir/test/python/global_constructors.py b/mlir/test/python/global_constructors.py
index 074ecd867ac7d..5020c00344a33 100644
--- a/mlir/test/python/global_constructors.py
+++ b/mlir/test/python/global_constructors.py
@@ -21,6 +21,7 @@ def run(f):
gc.collect()
assert Context._get_live_count() == 0
+
def lowerToLLVM(module):
pm = PassManager.parse(
"builtin.module(convert-func-to-llvm,reconcile-unrealized-casts)"
@@ -28,6 +29,7 @@ def lowerToLLVM(module):
pm.run(module.operation)
return module
+
# Test JIT callback in global constructor
# CHECK-LABEL: TEST: testJITCallbackInGlobalCtor
def testJITCallbackInGlobalCtor():
More information about the Mlir-commits
mailing list