[mlir][async] Make symbols in runtime library visible.

[ingomueller-net]

Summary

This patch exports the symbolds of the mlir_async_runtime library as
visible symbols (i.e., nm shows them as T instead of t). This is
necessary to be able to use load the library into the (Python binding
of) the execution engine, which, otherwise, can't see those symbols. The
symbols are made visible with the same mechanism as this is done in the
mlir_runner_utils and mlir_c_runner_utils (the main macro is copied from
there), which export the symbols for the same reason.

[mehdi_amini]

It shouldn't be needed, we expose a __mlir_runner_init symbol that register these functions explicitly with the JIT ExecutionEngine in JitRunner.cpp.

Instead of relying on symbol visibility, can we look into generalizing this kind of runtime library initialization?

[ingomueller-net]

To make sure I understand you and the current mechanisms correctly: the mechanism with __mlir_runner_init is currently used by mlir-cpu-runner for finding the exported symbols from the async runtime. This works even though the symbols in that library do not have public visibility. What you are proposing is (1) to have the execution engine load symbols through the same mechanism and potentially (2) to have other libraries (like the runner utils and the C runner utils) use the same mechanism (and revert their symbols back to private visibility (t)). Is that correct?

[mehdi_amini]

Yes: that seems more robust that relying on "polluting" the global namespace of the running binary with implicitly loading "all symbols from this shared library".
The JIT engine could just always look for a "magic" symbol to initialize a shared library on loading.

[ingomueller-net]

https://reviews.llvm.org/D153029 now contains a first stab at what I think you meant.

[ingomueller-net]

Abandoning in favor of https://reviews.llvm.org/D153029.

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

 #include <cstddef>
 #include <stdint.h>
 
 #ifdef mlir_async_runtime_EXPORTS
 // We are building this library
 #define MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS
 #endif // mlir_async_runtime_EXPORTS
 
+#ifdef _WIN32
+#ifndef MLIR_ASYNCRUNTIME_EXPORT
+#ifdef mlir_async_runtime_EXPORTS
+// We are building this library
+#define MLIR_ASYNCRUNTIME_EXPORT __declspec(dllexport)
+#else
+// We are using this library
+#define MLIR_ASYNCRUNTIME_EXPORT __declspec(dllimport)
+#endif // mlir_async_runtime_EXPORTS
+#endif // MLIR_ASYNCRUNTIME_EXPORT
+#else
+// Non-windows: use visibility attributes.
+#define MLIR_ASYNCRUNTIME_EXPORT __attribute__((visibility("default")))
+#endif // _WIN32
+
 namespace mlir {
 namespace runtime {
 
 //===----------------------------------------------------------------------===//
 // Async runtime API.
 //===----------------------------------------------------------------------===//
 
 // Runtime implementation of `async.token` data type.

 using CoroHandle = void *;           // coroutine handle
 using CoroResume = void (*)(void *); // coroutine resume function
 
 // Async runtime uses reference counting to manage the lifetime of async values
 // (values of async types like tokens, values and groups).
 using RefCountedObjPtr = void *;
 
 // Adds references to reference counted runtime object.
-extern "C" void mlirAsyncRuntimeAddRef(RefCountedObjPtr, int64_t);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+    mlirAsyncRuntimeAddRef(RefCountedObjPtr, int64_t);
 
 // Drops references from reference counted runtime object.
-extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr, int64_t);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+    mlirAsyncRuntimeDropRef(RefCountedObjPtr, int64_t);
 
 // Create a new `async.token` in not-ready state.
-extern "C" AsyncToken *mlirAsyncRuntimeCreateToken();
+extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncToken *mlirAsyncRuntimeCreateToken();
 
 // Create a new `async.value` in not-ready state. Size parameter specifies the
 // number of bytes that will be allocated for the async value storage. Storage
 // is owned by the `async.value` and deallocated when the async value is
 // destructed (reference count drops to zero).
-extern "C" AsyncValue *mlirAsyncRuntimeCreateValue(int64_t);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncValue *
+    mlirAsyncRuntimeCreateValue(int64_t);
 
 // Create a new `async.group` in empty state.
-extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup(int64_t size);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT AsyncGroup *
+mlirAsyncRuntimeCreateGroup(int64_t size);
 
-extern "C" int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT int64_t
+mlirAsyncRuntimeAddTokenToGroup(AsyncToken *, AsyncGroup *);
 
 // Switches `async.token` to ready state and runs all awaiters.
-extern "C" void mlirAsyncRuntimeEmplaceToken(AsyncToken *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeEmplaceToken(AsyncToken *);
 
 // Switches `async.value` to ready state and runs all awaiters.
-extern "C" void mlirAsyncRuntimeEmplaceValue(AsyncValue *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeEmplaceValue(AsyncValue *);
 
 // Switches `async.token` to error state and runs all awaiters.
-extern "C" void mlirAsyncRuntimeSetTokenError(AsyncToken *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeSetTokenError(AsyncToken *);
 
 // Switches `async.value` to error state and runs all awaiters.
-extern "C" void mlirAsyncRuntimeSetValueError(AsyncValue *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeSetValueError(AsyncValue *);
 
 // Returns true if token is in the error state.
-extern "C" bool mlirAsyncRuntimeIsTokenError(AsyncToken *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT bool
+mlirAsyncRuntimeIsTokenError(AsyncToken *);
 
 // Returns true if value is in the error state.
-extern "C" bool mlirAsyncRuntimeIsValueError(AsyncValue *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT bool
+mlirAsyncRuntimeIsValueError(AsyncValue *);
 
 // Returns true if group is in the error state (any of the tokens or values
 // added to the group are in the error state).
-extern "C" bool mlirAsyncRuntimeIsGroupError(AsyncGroup *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT bool
+mlirAsyncRuntimeIsGroupError(AsyncGroup *);
 
 // Blocks the caller thread until the token becomes ready.
-extern "C" void mlirAsyncRuntimeAwaitToken(AsyncToken *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitToken(AsyncToken *);
 
 // Blocks the caller thread until the value becomes ready.
-extern "C" void mlirAsyncRuntimeAwaitValue(AsyncValue *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitValue(AsyncValue *);
 
 // Blocks the caller thread until the elements in the group become ready.
-extern "C" void mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *);
 
 // Returns a pointer to the storage owned by the async value.
-extern "C" ValueStorage mlirAsyncRuntimeGetValueStorage(AsyncValue *);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT ValueStorage
+mlirAsyncRuntimeGetValueStorage(AsyncValue *);
 
 // Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime.
-extern "C" void mlirAsyncRuntimeExecute(CoroHandle, CoroResume);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimeExecute(CoroHandle,
+                                                                 CoroResume);
 
 // Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime after the token becomes ready.
-extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *, CoroHandle,
-                                                     CoroResume);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *, CoroHandle, CoroResume);
 
 // Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime after the value becomes ready.
-extern "C" void mlirAsyncRuntimeAwaitValueAndExecute(AsyncValue *, CoroHandle,
-                                                     CoroResume);
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
+mlirAsyncRuntimeAwaitValueAndExecute(AsyncValue *, CoroHandle, CoroResume);
 
 // Executes the task (coro handle + resume function) in one of the threads
 // managed by the runtime after the all members of the group become ready.
-extern "C" void
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void
 mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume);
 
 // Returns the current number of available worker threads in the threadpool.
-extern "C" int64_t mlirAsyncRuntimGetNumWorkerThreads();
+extern "C" MLIR_ASYNCRUNTIME_EXPORT int64_t
+mlirAsyncRuntimGetNumWorkerThreads();
 
 //===----------------------------------------------------------------------===//
 // Small async runtime support library for testing.
 //===----------------------------------------------------------------------===//
 
-extern "C" void mlirAsyncRuntimePrintCurrentThreadId();
+extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimePrintCurrentThreadId();
 
 } // namespace runtime
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_