[mlir] AsyncRuntime: mode runtime declarations to mlir::runtime namespace

[ezhulenev]

Summary

Define Async runtime related typedefs in the mlir::runtime namespace.

[ezhulenev]

Rebase

[ezhulenev]

Keep extern "C" in API function definition

[mehdi_amini]

Thanks!

mlir/include/mlir/ExecutionEngine/AsyncRuntime.h

 #define MLIR_ASYNCRUNTIME_EXPORT __declspec(dllimport)
 #endif // mlir_async_runtime_EXPORTS
 #endif // MLIR_ASYNCRUNTIME_EXPORT
 #else
 #define MLIR_ASYNCRUNTIME_EXPORT
 #define MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS
 #endif // _WIN32
 
+namespace mlir {
+namespace runtime {
+
 //===----------------------------------------------------------------------===//
 // Async runtime API.
 //===----------------------------------------------------------------------===//
 
 // Runtime implementation of `async.token` data type.
 typedef struct AsyncToken AsyncToken;
 
 // Runtime implementation of `async.group` data type.

 mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *, CoroHandle, CoroResume);
 
 //===----------------------------------------------------------------------===//
 // Small async runtime support library for testing.
 //===----------------------------------------------------------------------===//
 
 extern "C" MLIR_ASYNCRUNTIME_EXPORT void mlirAsyncRuntimePrintCurrentThreadId();
 
+} // namespace runtime
+} // namespace mlir
+
 #endif // MLIR_EXECUTIONENGINE_ASYNCRUNTIME_H_

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

 #include <cassert>
 #include <condition_variable>
 #include <functional>
 #include <iostream>
 #include <mutex>
 #include <thread>
 #include <vector>
 
+using namespace mlir::runtime;
+
 //===----------------------------------------------------------------------===//
 // Async runtime API.
 //===----------------------------------------------------------------------===//
 
+namespace mlir {
+namespace runtime {
 namespace {
 
 // Forward declare class defined below.
 class RefCounted;
 
 // -------------------------------------------------------------------------- //
 // AsyncRuntime orchestrates all async operations and Async runtime API is built
 // on top of the default runtime instance.
private:
   }
   void dropNumRefCountedObjects() {
     numRefCountedObjects.fetch_sub(1, std::memory_order_relaxed);
   }
 
   std::atomic<int32_t> numRefCountedObjects;
 };
 
-// Returns the default per-process instance of an async runtime.
-AsyncRuntime *getDefaultAsyncRuntimeInstance() {
-  static auto runtime = std::make_unique<AsyncRuntime>();
-  return runtime.get();
-}
-
 // -------------------------------------------------------------------------- //
 // A base class for all reference counted objects created by the async runtime.
 // -------------------------------------------------------------------------- //
 
 class RefCounted {
 public:
   RefCounted(AsyncRuntime *runtime, int32_t refCount = 1)
       : runtime(runtime), refCount(refCount) {

 
 private:
   AsyncRuntime *runtime;
   std::atomic<int32_t> refCount;
 };
 
 } // namespace
 
+// Returns the default per-process instance of an async runtime.
+static AsyncRuntime *getDefaultAsyncRuntimeInstance() {
+  static auto runtime = std::make_unique<AsyncRuntime>();
+  return runtime.get();
+}
+
 struct AsyncToken : public RefCounted {
   // AsyncToken created with a reference count of 2 because it will be returned
   // to the `async.execute` caller and also will be later on emplaced by the
   // asynchronously executed task. If the caller immediately will drop its
   // reference we must ensure that the token will be alive until the
   // asynchronous operation is completed.
   AsyncToken(AsyncRuntime *runtime) : RefCounted(runtime, /*count=*/2) {}
 
struct AsyncGroup : public RefCounted {
 
   // Internal state below guarded by a mutex.
   std::mutex mu;
   std::condition_variable cv;
 
   std::vector<std::function<void()>> awaiters;
 };
 
+} // namespace runtime
+} // namespace mlir
+
 // Adds references to reference counted runtime object.
-extern "C" void mlirAsyncRuntimeAddRef(RefCountedObjPtr ptr, int32_t count) {
+void mlirAsyncRuntimeAddRef(RefCountedObjPtr ptr, int32_t count) {
   RefCounted *refCounted = static_cast<RefCounted *>(ptr);
   refCounted->addRef(count);
 }
 
 // Drops references from reference counted runtime object.
-extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr ptr, int32_t count) {
+void mlirAsyncRuntimeDropRef(RefCountedObjPtr ptr, int32_t count) {
   RefCounted *refCounted = static_cast<RefCounted *>(ptr);
   refCounted->dropRef(count);
 }
 
 // Create a new `async.token` in not-ready state.
-extern "C" AsyncToken *mlirAsyncRuntimeCreateToken() {
+AsyncToken *mlirAsyncRuntimeCreateToken() {
   AsyncToken *token = new AsyncToken(getDefaultAsyncRuntimeInstance());
   return token;
 }
 
 // Create a new `async.group` in empty state.
-extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup() {
+AsyncGroup *mlirAsyncRuntimeCreateGroup() {
   AsyncGroup *group = new AsyncGroup(getDefaultAsyncRuntimeInstance());
   return group;
 }
 
-extern "C" int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *token,
-                                                   AsyncGroup *group) {
+int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *token, AsyncGroup *group) {
   std::unique_lock<std::mutex> lockToken(token->mu);
   std::unique_lock<std::mutex> lockGroup(group->mu);
 
   // Get the rank of the token inside the group before we drop the reference.
   int rank = group->rank.fetch_add(1);
   group->pendingTokens.fetch_add(1);
 
   auto onTokenReady = [group]() {
token->awaiters.push_back([group, onTokenReady]() {
       group->dropRef();
     });
   }
 
   return rank;
 }
 
 // Switches `async.token` to ready state and runs all awaiters.
-extern "C" void mlirAsyncRuntimeEmplaceToken(AsyncToken *token) {
+void mlirAsyncRuntimeEmplaceToken(AsyncToken *token) {
   // Make sure that `dropRef` does not destroy the mutex owned by the lock.
   {
     std::unique_lock<std::mutex> lock(token->mu);
     token->ready = true;
     token->cv.notify_all();
     for (auto &awaiter : token->awaiters)
       awaiter();
   }
 
   // Async tokens created with a ref count `2` to keep token alive until the
   // async task completes. Drop this reference explicitly when token emplaced.
   token->dropRef();
 }
 
-extern "C" void mlirAsyncRuntimeAwaitToken(AsyncToken *token) {
+void mlirAsyncRuntimeAwaitToken(AsyncToken *token) {
   std::unique_lock<std::mutex> lock(token->mu);
   if (!token->ready)
     token->cv.wait(lock, [token] { return token->ready; });
 }
 
-extern "C" void mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *group) {
+void mlirAsyncRuntimeAwaitAllInGroup(AsyncGroup *group) {
   std::unique_lock<std::mutex> lock(group->mu);
   if (group->pendingTokens != 0)
     group->cv.wait(lock, [group] { return group->pendingTokens == 0; });
 }
 
-extern "C" void mlirAsyncRuntimeExecute(CoroHandle handle, CoroResume resume) {
+void mlirAsyncRuntimeExecute(CoroHandle handle, CoroResume resume) {
   (*resume)(handle);
 }
 
-extern "C" void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token,
-                                                     CoroHandle handle,
-                                                     CoroResume resume) {
+void mlirAsyncRuntimeAwaitTokenAndExecute(AsyncToken *token, CoroHandle handle,
   std::unique_lock<std::mutex> lock(token->mu);
   auto execute = [handle, resume]() { (*resume)(handle); };
   if (token->ready)
     execute();
   else
     token->awaiters.push_back([execute]() { execute(); });
 }
 
-extern "C" void mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *group,
-                                                          CoroHandle handle,
-                                                          CoroResume resume) {
+void mlirAsyncRuntimeAwaitAllInGroupAndExecute(AsyncGroup *group,
   std::unique_lock<std::mutex> lock(group->mu);
   auto execute = [handle, resume]() { (*resume)(handle); };
   if (group->pendingTokens == 0)
     execute();
   else
     group->awaiters.push_back([execute]() { execute(); });
 }
 
 //===----------------------------------------------------------------------===//
 // Small async runtime support library for testing.
 //===----------------------------------------------------------------------===//
 
-extern "C" void mlirAsyncRuntimePrintCurrentThreadId() {
+void mlirAsyncRuntimePrintCurrentThreadId() {
   static thread_local std::thread::id thisId = std::this_thread::get_id();
   std::cout << "Current thread id: " << thisId << std::endl;
 }
 
 #endif // MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS