[Mlir-commits] [mlir] Reland [mlir][ExecutionEngine] Add support for global constructors and destructors #78070 (PR #78170)

[Fabian Mora]

https://github.com/fabianmcg created https://github.com/llvm/llvm-project/pull/78170

This patch add support for executing global constructors and destructors in the ExecutionEngine.

>From a72a40a139271d1b8174e109a56fc201a64b3a46 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Sat, 13 Jan 2024 22:31:42 +0000
Subject: [PATCH 1/2] [mlir][ExecutionEngine] Add support for global
 constructors and destructors

This patch add support for executing global constructors and destructors in the
`ExecutionEngine`.
---
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp  |  6 ++++
 .../mlir-cpu-runner/global-constructors.mlir  | 32 +++++++++++++++++++
 2 files changed, 38 insertions(+)
 create mode 100644 mlir/test/mlir-cpu-runner/global-constructors.mlir

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index 267ec236d3fd590..d4ad2da045e2c4a 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -219,6 +219,9 @@ ExecutionEngine::ExecutionEngine(bool enableObjectDump,
 }
 
 ExecutionEngine::~ExecutionEngine() {
+  // Execute the global destructors from the module being processed.
+  if (jit)
+    llvm::consumeError(jit->deinitialize(jit->getMainJITDylib()));
   // Run all dynamic library destroy callbacks to prepare for the shutdown.
   for (LibraryDestroyFn destroy : destroyFns)
     destroy();
@@ -396,6 +399,9 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
   };
   engine->registerSymbols(runtimeSymbolMap);
 
+  // Execute the global constructors from the module being processed.
+  cantFail(engine->jit->initialize(engine->jit->getMainJITDylib()));
+
   return std::move(engine);
 }
 
diff --git a/mlir/test/mlir-cpu-runner/global-constructors.mlir b/mlir/test/mlir-cpu-runner/global-constructors.mlir
new file mode 100644
index 000000000000000..3ebcaeccc83e32b
--- /dev/null
+++ b/mlir/test/mlir-cpu-runner/global-constructors.mlir
@@ -0,0 +1,32 @@
+// RUN: mlir-cpu-runner %s -e entry -entry-point-result=void  \
+// RUN: -shared-libs=%mlir_c_runner_utils | \
+// RUN: FileCheck %s
+
+// Test that the `ctor` executes before `entry` and that `dtor` executes last.
+module {
+  llvm.func @printNewline()
+  llvm.func @printI64(i64)
+  llvm.mlir.global_ctors {ctors = [@ctor], priorities = [0 : i32]}
+  llvm.mlir.global_dtors {dtors = [@dtor], priorities = [0 : i32]}
+  llvm.func @ctor() {
+    %0 = llvm.mlir.constant(1 : i64) : i64
+    llvm.call @printI64(%0) : (i64) -> ()
+    llvm.call @printNewline() : () -> ()
+    // CHECK: 1
+    llvm.return
+  }
+  llvm.func @entry() {
+    %0 = llvm.mlir.constant(2 : i64) : i64
+    llvm.call @printI64(%0) : (i64) -> ()
+    llvm.call @printNewline() : () -> ()
+    // CHECK: 2
+    llvm.return
+  }
+  llvm.func @dtor() {
+    %0 = llvm.mlir.constant(3 : i64) : i64
+    llvm.call @printI64(%0) : (i64) -> ()
+    llvm.call @printNewline() : () -> ()
+    // CHECK: 3
+    llvm.return
+  }
+}

>From 2877cc1e90d57c4db058e26014dcc53c9b7097e9 Mon Sep 17 00:00:00 2001
From: Fabian Mora <fmora.dev at gmail.com>
Date: Mon, 15 Jan 2024 15:32:55 +0000
Subject: [PATCH 2/2] add unsupported AArch64

---
 mlir/lib/ExecutionEngine/ExecutionEngine.cpp       | 9 +++++++--
 mlir/test/mlir-cpu-runner/global-constructors.mlir | 1 +
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
index d4ad2da045e2c4a..65f674d25ece79a 100644
--- a/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -220,7 +220,9 @@ ExecutionEngine::ExecutionEngine(bool enableObjectDump,
 
 ExecutionEngine::~ExecutionEngine() {
   // Execute the global destructors from the module being processed.
-  if (jit)
+  // TODO: Allow JIT initialize for AArch64. Currently there's a bug causing a
+  // crash for AArch64 see related issue #71963.
+  if (jit && !jit->getTargetTriple().isAArch64())
     llvm::consumeError(jit->deinitialize(jit->getMainJITDylib()));
   // Run all dynamic library destroy callbacks to prepare for the shutdown.
   for (LibraryDestroyFn destroy : destroyFns)
@@ -400,7 +402,10 @@ ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
   engine->registerSymbols(runtimeSymbolMap);
 
   // Execute the global constructors from the module being processed.
-  cantFail(engine->jit->initialize(engine->jit->getMainJITDylib()));
+  // 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);
 }
diff --git a/mlir/test/mlir-cpu-runner/global-constructors.mlir b/mlir/test/mlir-cpu-runner/global-constructors.mlir
index 3ebcaeccc83e32b..fe1fe06574648e7 100644
--- a/mlir/test/mlir-cpu-runner/global-constructors.mlir
+++ b/mlir/test/mlir-cpu-runner/global-constructors.mlir
@@ -1,3 +1,4 @@
+// UNSUPPORTED: target=aarch64{{.*}}
 // RUN: mlir-cpu-runner %s -e entry -entry-point-result=void  \
 // RUN: -shared-libs=%mlir_c_runner_utils | \
 // RUN: FileCheck %s