[mlir][RunnerUtils] Make symbols private + implement loading mechanism.

[ingomueller-net]

Summary

There are two ways to make symbols from a shared library visible in the
execution engine: exporting the symbols with public visibility or
implementing a loading/unloading mechansim that registers the exported
symbols explicitly. The latter has only been available in the JIT runner
until recently, but https://reviews.llvm.org/D153029 makes it available
in any usage of the execution engine (including the Python bindings).

This patch makes the runner utils library use the latter mechanism
instead of the former, i.e., it makes all of its symbols private and
implements the init/destroy functions of the loading mechanism to
control explicitly which symbols it registers.

[ingomueller-net]

(This still fails in CI until the issues discussed in https://reviews.llvm.org/D153029 have been resolved.)

[ingomueller-net]

I am not sure whether this is actually a good idea. It makes it impossible to use the library with dlopen. Some of the integration tests of the sparse compile actually do that: they run their code through lli --dlopen=%mlir_runner_utils, so the loading mechanism of the execution engine isn't involved and the symbols don't get loaded/aren't visible to the code. Not sure how to proceed now.

[ingomueller-net]

(Rebasing on top of https://reviews.llvm.org/D153255 to have the visibility issues fixed.)

[ingomueller-net]

(Rabasing onto main. I believe my previous rebase didn't result in https://reviews.llvm.org/D153255 being included, so CI failed.)

[mehdi_amini]

It makes it impossible to use the library with dlopen

I also makes it impossible to use in AOT mode by directly linking to it I think?

[ingomueller-net]

In D153250#4432425, @mehdi_amini wrote:

It makes it impossible to use the library with dlopen

I also makes it impossible to use in AOT mode by directly linking to it I think?

I thinks so too, yes.

[ingomueller-net]

In D153250#4432461, @ingomueller-net wrote:

In D153250#4432425, @mehdi_amini wrote:

It makes it impossible to use the library with dlopen

I also makes it impossible to use in AOT mode by directly linking to it I think?

I thinks so too, yes.

This makes me wonder whether the loading mechanism is such a good idea after all. I looked at D94270/D94312, which introduced them, and couldn't quite figure out why they were introduced. @ezhulenev: can you shed some light?

[mehdi_amini]

I don't quite get what is the problem you're trying to solve here actually?

[ingomueller-net]

In D153250#4434366, @mehdi_amini wrote:

I don't quite get what is the problem you're trying to solve here actually?

I don't understand why the async runtime lib uses one loading mechanism (explicit init/destroy functions) and the runner utils lib uses another (import all exported symbols via ORC). I thought that the init/destroy mechanism was preferred (hence https://reviews.llvm.org/D153009 was rejected), so I am trying to make things consistent with this patch by using the same mechanism for the runner lib. If that mechanism has problems, though, don't they apply for the async lib as well? If so, shouldn't we use the other mechanism there as well?

[mehdi_amini]

Why should the explicit loading mechanism imply to make all symbol private though?

[ingomueller-net]

In D153250#4434438, @mehdi_amini wrote:

Why should the explicit loading mechanism imply to make all symbol private though?

Ha, yes! I didn't see the obvious... Thanks for helping out ;) I'll prepare a new version...

[ingomueller-net]

Rebasing to main and making the symbols of the library public again (but keeping all unrelated symbols hidden).

[mehdi_amini]

mlir/lib/ExecutionEngine/CMakeLists.txt

167

Are these CMake directives needed?

[mehdi_amini]

mlir/lib/ExecutionEngine/CMakeLists.txt

167

(Nevermind this comment)

[hctim]

Hi folks, I did a bisect and unfortunately this broke the HWASan buildbot. It only showed up on the build after yours because your patch was included in a set that was broken for other reasons. https://lab.llvm.org/buildbot/#/builders/236/builds/4911

I'll check tomorrow and see exactly why your patch is causing a problem, but dynamic loading can play interesting with HWAddressSanitizer, because we need to ensure that global variables have the correct memory tag applied, as the address tag is synthesized in the symbols in the DSO.

If you have access to an aarch64 machine and want to reproduce, please feel free. The fastest way is to use the instructions at https://github.com/google/sanitizers/wiki/SanitizerBotReproduceBuild, the script that's of note here is zorg/buildbot/builders/sanitizers/buildbot_bootstrap_hwasan.sh. You may find commenting out check_stage1_hwasan, build_stage3_hwasan, and check_stage3_hwasan, as well as only running the MLIR tests in stage2 helpful:

diff --git a/zorg/buildbot/builders/sanitizers/buildbot_functions.sh b/zorg/buildbot/builders/sanitizers/buildbot_functions.sh
index 1b64a929..e2f96a09 100755
--- a/zorg/buildbot/builders/sanitizers/buildbot_functions.sh
+++ b/zorg/buildbot/builders/sanitizers/buildbot_functions.sh
@@ -388,7 +388,7 @@ function check_stage2 {
   fi

   echo @@@BUILD_STEP stage2/$sanitizer_name check@@@
-  ninja -C ${STAGE2_DIR} check-all || build_failure
+  ninja -C ${STAGE2_DIR} check-mlir || build_failure

   if [[ "${STAGE2_SKIP_TEST_CXX:-}" != "1" ]] ; then
     echo @@@BUILD_STEP stage2/$sanitizer_name check-cxx@@@

[ingomueller-net]

@hctim, @mehdi_amini: Thanks for tracking down the problem and for reverting! Sorry about these issues. After looking around a bit, I decided that it is too much hassle for me to reproduce the issue in order to fix it, and currently too unimportant to spend more time. I'll leave the commit reverted as is for the time being...

mlir/lib/ExecutionEngine/CMakeLists.txt

if(LLVM_ENABLE_PIC)
   target_compile_definitions(mlir_c_runner_utils PRIVATE mlir_c_runner_utils_EXPORTS)
 
   add_mlir_library(mlir_runner_utils
     SHARED
     RunnerUtils.cpp
 
     EXCLUDE_FROM_LIBMLIR
   )
+  set_property(TARGET mlir_runner_utils PROPERTY CXX_VISIBILITY_PRESET hidden)
   target_compile_definitions(mlir_runner_utils PRIVATE mlir_runner_utils_EXPORTS)
+  if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    # Don't export symbols from link-time dependencies, these are internal
+    # implementation details.
+    # FIXME: Add a similar fix for Windows.
+    target_link_options(mlir_runner_utils PRIVATE "-Wl,-exclude-libs,ALL")
+  endif()
 
   add_mlir_library(mlir_async_runtime
     SHARED
     AsyncRuntime.cpp
 
     EXCLUDE_FROM_LIBMLIR
 
     LINK_LIBS PUBLIC
     ${LLVM_PTHREAD_LIB}

mlir/lib/ExecutionEngine/RunnerUtils.cpp

 //===- RunnerUtils.cpp - Utils for MLIR exec on targets with a C++ runtime ===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements basic functions to debug structured MLIR types at
 // runtime. Entities in this file may not be compatible with targets without a
 // C++ runtime. These may be progressively migrated to CRunnerUtils.cpp over
 // time.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/RunnerUtils.h"
+#include "llvm/ADT/StringMap.h"
 #include <chrono>
 
 // NOLINTBEGIN(*-identifier-naming)
 
 extern "C" void _mlir_ciface_printMemrefShapeI8(UnrankedMemRefType<int8_t> *M) {
   std::cout << "Unranked Memref ";
   printMemRefMetaData(std::cout, DynamicMemRefType<int8_t>(*M));
   std::cout << "\n";

 
 extern "C" int64_t verifyMemRefC64(int64_t rank, void *actualPtr,
                                    void *expectedPtr) {
   UnrankedMemRefType<impl::complex64> actualDesc = {rank, actualPtr};
   UnrankedMemRefType<impl::complex64> expectedDesc = {rank, expectedPtr};
   return _mlir_ciface_verifyMemRefC64(&actualDesc, &expectedDesc);
 }
 
+//===----------------------------------------------------------------------===//
+// MLIR ExecutionEngine dynamic library integration.
+//===----------------------------------------------------------------------===//
+
+// Visual Studio had a bug that fails to compile nested generic lambdas
+// inside an `extern "C"` function.
+//   https://developercommunity.visualstudio.com/content/problem/475494/clexe-error-with-lambda-inside-function-templates.html
+// The bug is fixed in VS2019 16.1. Separating the declaration and definition is
+// a work around for older versions of Visual Studio.
+extern "C" MLIR_RUNNERUTILS_EXPORT void
+__mlir_execution_engine_init(llvm::StringMap<void *> &exportSymbols);
+
+void __mlir_execution_engine_init(llvm::StringMap<void *> &exportSymbols) {
+  auto exportSymbol = [&](llvm::StringRef name, auto ptr) {
+    assert(exportSymbols.count(name) == 0 && "symbol already exists");
+    exportSymbols[name] = reinterpret_cast<void *>(ptr);
+  };
+
+  exportSymbol("_mlir_ciface_printMemrefShapeI8",
+               &_mlir_ciface_printMemrefShapeI8);
+  exportSymbol("_mlir_ciface_printMemrefShapeI32",
+               &_mlir_ciface_printMemrefShapeI32);
+  exportSymbol("_mlir_ciface_printMemrefShapeI64",
+               &_mlir_ciface_printMemrefShapeI64);
+  exportSymbol("_mlir_ciface_printMemrefShapeF32",
+               &_mlir_ciface_printMemrefShapeF32);
+  exportSymbol("_mlir_ciface_printMemrefShapeF64",
+               &_mlir_ciface_printMemrefShapeF64);
+  exportSymbol("_mlir_ciface_printMemrefShapeInd",
+               &_mlir_ciface_printMemrefShapeInd);
+  exportSymbol("_mlir_ciface_printMemrefShapeC32",
+               &_mlir_ciface_printMemrefShapeC32);
+  exportSymbol("_mlir_ciface_printMemrefShapeC64",
+               &_mlir_ciface_printMemrefShapeC64);
+  exportSymbol("_mlir_ciface_printMemrefI8", &_mlir_ciface_printMemrefI8);
+  exportSymbol("_mlir_ciface_printMemrefI32", &_mlir_ciface_printMemrefI32);
+  exportSymbol("_mlir_ciface_printMemrefI64", &_mlir_ciface_printMemrefI64);
+  exportSymbol("_mlir_ciface_printMemrefF32", &_mlir_ciface_printMemrefF32);
+  exportSymbol("_mlir_ciface_printMemrefF64", &_mlir_ciface_printMemrefF64);
+  exportSymbol("_mlir_ciface_printMemrefInd", &_mlir_ciface_printMemrefInd);
+  exportSymbol("_mlir_ciface_printMemrefC32", &_mlir_ciface_printMemrefC32);
+  exportSymbol("_mlir_ciface_printMemrefC64", &_mlir_ciface_printMemrefC64);
+  exportSymbol("_mlir_ciface_nanoTime", &_mlir_ciface_nanoTime);
+  exportSymbol("printMemrefI32", &printMemrefI32);
+  exportSymbol("printMemrefI64", &printMemrefI64);
+  exportSymbol("printMemrefF32", &printMemrefF32);
+  exportSymbol("printMemrefF64", &printMemrefF64);
+  exportSymbol("printMemrefInd", &printMemrefInd);
+  exportSymbol("printMemrefC32", &printMemrefC32);
+  exportSymbol("printMemrefC64", &printMemrefC64);
+  exportSymbol("printCString", &printCString);
+  exportSymbol("_mlir_ciface_printMemref0dF32", &_mlir_ciface_printMemref0dF32);
+  exportSymbol("_mlir_ciface_printMemref1dF32", &_mlir_ciface_printMemref1dF32);
+  exportSymbol("_mlir_ciface_printMemref2dF32", &_mlir_ciface_printMemref2dF32);
+  exportSymbol("_mlir_ciface_printMemref3dF32", &_mlir_ciface_printMemref3dF32);
+  exportSymbol("_mlir_ciface_printMemref4dF32", &_mlir_ciface_printMemref4dF32);
+  exportSymbol("_mlir_ciface_printMemref1dI8", &_mlir_ciface_printMemref1dI8);
+  exportSymbol("_mlir_ciface_printMemref1dI32", &_mlir_ciface_printMemref1dI32);
+  exportSymbol("_mlir_ciface_printMemref1dI64", &_mlir_ciface_printMemref1dI64);
+  exportSymbol("_mlir_ciface_printMemref1dF64", &_mlir_ciface_printMemref1dF64);
+  exportSymbol("_mlir_ciface_printMemref1dInd", &_mlir_ciface_printMemref1dInd);
+  exportSymbol("_mlir_ciface_printMemref1dC32", &_mlir_ciface_printMemref1dC32);
+  exportSymbol("_mlir_ciface_printMemref1dC64", &_mlir_ciface_printMemref1dC64);
+  exportSymbol("_mlir_ciface_printMemrefVector4x4xf32",
+               &_mlir_ciface_printMemrefVector4x4xf32);
+  exportSymbol("_mlir_ciface_verifyMemRefI32", &_mlir_ciface_verifyMemRefI32);
+  exportSymbol("_mlir_ciface_verifyMemRefF32", &_mlir_ciface_verifyMemRefF32);
+  exportSymbol("_mlir_ciface_verifyMemRefF64", &_mlir_ciface_verifyMemRefF64);
+  exportSymbol("_mlir_ciface_verifyMemRefInd", &_mlir_ciface_verifyMemRefInd);
+  exportSymbol("_mlir_ciface_verifyMemRefC32", &_mlir_ciface_verifyMemRefC32);
+  exportSymbol("_mlir_ciface_verifyMemRefC64", &_mlir_ciface_verifyMemRefC64);
+  exportSymbol("verifyMemRefI32", &verifyMemRefI32);
+  exportSymbol("verifyMemRefF32", &verifyMemRefF32);
+  exportSymbol("verifyMemRefF64", &verifyMemRefF64);
+  exportSymbol("verifyMemRefInd", &verifyMemRefInd);
+  exportSymbol("verifyMemRefC32", &verifyMemRefC32);
+  exportSymbol("verifyMemRefC64", &verifyMemRefC64);
+}
+
+extern "C" MLIR_RUNNERUTILS_EXPORT void __mlir_execution_engine_destroy() {}
+
 // NOLINTEND(*-identifier-naming)