[mlir][CAPI] Expose the rest of MLIRContext's constructors

[krzysz00]

Summary

It's recommended practice that people calling MLIR in a loop
pre-create a LLVM ThreadPool and a dialect registry and then
explicitly pass those into a MLIRContext for each compilation.
However, the C API does not expose the functions needed to follow this
recommendation from a project that isn't calling MLIR's C++ dilectly.

Add the necessary APIs to mlir-c, including a wrapper around LLVM's
ThreadPool struct (so as to avoid having to amend or re-export parts
of the LLVM API).

[makslevental]

not sure about the connection to DialectRegistry but can't you just call mlirContextEnableMultithreading which will do impl->ownedThreadPool = std::make_unique<llvm::ThreadPool>(); (and thus similarly create a hardware_concurrency() threadpool)?

https://github.com/llvm/llvm-project/blob/9566ee280607d91fa2e5eca730a6765ac84dfd0f/mlir/lib/CAPI/IR/IR.cpp#L79

Not that I'm opposed to this change or anything but just double-checking my own understanding.

[krzysz00]

@makslevental The purpose of these constructors is to specifically not create that owned thread pool.

The intended usage can be seen in https://github.com/ROCmSoftwarePlatform/AMDMIGraphX/blob/develop/src/targets/gpu/mlir.cpp#L178-L186 , which is a class whose instances will be run in parallel. We don't want to create one thread pool for each kernel compilation, but instead share a thread pool across all those instances.

[makslevental]

In D153593#4478108, @krzysz00 wrote:

@makslevental The purpose of these constructors is to specifically not create that owned thread pool.

The intended usage can be seen in https://github.com/ROCmSoftwarePlatform/AMDMIGraphX/blob/develop/src/targets/gpu/mlir.cpp#L178-L186 , which is a class whose instances will be run in parallel. We don't want to create one thread pool for each kernel compilation, but instead share a thread pool across all those instances.

gotcha - not sure why it wasn't clear from the beginning - thanks for spelling it out for me.

mlir/include/mlir-c/IR.h

 };
 typedef struct MlirNamedAttribute MlirNamedAttribute;
 
 //===----------------------------------------------------------------------===//
 // Context API.
 //===----------------------------------------------------------------------===//
 
 /// Creates an MLIR context and transfers its ownership to the caller.
+/// This sets the default multithreading option (enabled).
 MLIR_CAPI_EXPORTED MlirContext mlirContextCreate(void);
 
+/// Creates an MLIR context with an explicit setting of the multithreading
+/// setting and transfers its ownership to the caller.
+MLIR_CAPI_EXPORTED MlirContext
+mlirContextCreateWithThreading(bool threadingEnabled);
+
+/// Creates an MLIR context, setting the multithreading setting explicitly and
+/// pre-loading the dialects from the provided DialectRegistry.
+MLIR_CAPI_EXPORTED MlirContext mlirContextCreateWithRegistry(
+    MlirDialectRegistry registry, bool threadingEnabled);
+
 /// Checks if two contexts are equal.
 MLIR_CAPI_EXPORTED bool mlirContextEqual(MlirContext ctx1, MlirContext ctx2);
 
 /// Checks whether a context is null.
 static inline bool mlirContextIsNull(MlirContext context) {
   return !context.ptr;
 }
 

 
 /// Returns whether the given fully-qualified operation (i.e.
 /// 'dialect.operation') is registered with the context. This will return true
 /// if the dialect is loaded and the operation is registered within the
 /// dialect.
 MLIR_CAPI_EXPORTED bool mlirContextIsRegisteredOperation(MlirContext context,
                                                          MlirStringRef name);
 
+/// Sets the thread pool of the context explicitly, enabling multithreading in
+/// the process. This API should be used to avoid re-creating thread pools in
+/// long-running applications that perform multiple compilations, see
+/// the C++ documentation for MLIRContext for details.
+MLIR_CAPI_EXPORTED void mlirContextSetThreadPool(MlirContext context,
+                                                 MlirLlvmThreadPool threadPool);
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//
 
 /// Returns the context that owns the dialect.
 MLIR_CAPI_EXPORTED MlirContext mlirDialectGetContext(MlirDialect dialect);
 
 /// Checks if the dialect is null.

mlir/include/mlir-c/Support.h

 #endif
 
 #define DEFINE_C_API_STRUCT(name, storage)                                     \
   struct name {                                                                \
     storage *ptr;                                                              \
   };                                                                           \
   typedef struct name name
 
+/// Re-export llvm::ThreadPool so as to avoid including the LLVM C API directly.
+DEFINE_C_API_STRUCT(MlirLlvmThreadPool, void);
 DEFINE_C_API_STRUCT(MlirTypeID, const void);
 DEFINE_C_API_STRUCT(MlirTypeIDAllocator, void);
 
 #undef DEFINE_C_API_STRUCT
 
 //===----------------------------------------------------------------------===//
 // MlirStringRef.
 //===----------------------------------------------------------------------===//

 
 /// Creates a logical result representing a failure.
 inline static MlirLogicalResult mlirLogicalResultFailure(void) {
   MlirLogicalResult res = {0};
   return res;
 }
 
 //===----------------------------------------------------------------------===//
+// MlirLlvmThreadPool.
+//===----------------------------------------------------------------------===//
+
+/// Create an LLVM thread pool. This is reexported here to avoid directly
+/// pulling in the LLVM headers directly.
+MLIR_CAPI_EXPORTED MlirLlvmThreadPool mlirLlvmThreadPoolCreate(void);
+
+/// Destroy an LLVM thread pool.
+MLIR_CAPI_EXPORTED void mlirLlvmThreadPoolDestroy(MlirLlvmThreadPool pool);
+
+//===----------------------------------------------------------------------===//
 // TypeID API.
 //===----------------------------------------------------------------------===//
 
 /// `ptr` must be 8 byte aligned and unique to a type valid for the duration of
 /// the returned type id's usage
 MLIR_CAPI_EXPORTED MlirTypeID mlirTypeIDCreate(const void *ptr);
 
 /// Checks whether a type id is null.

mlir/include/mlir/CAPI/Support.h

 #define MLIR_CAPI_SUPPORT_H
 
 #include "mlir-c/Support.h"
 #include "mlir/CAPI/Wrap.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/StringRef.h"
 
+namespace llvm {
+class ThreadPool;
+} // namespace llvm
+
 /// Converts a StringRef into its MLIR C API equivalent.
 inline MlirStringRef wrap(llvm::StringRef ref) {
   return mlirStringRefCreate(ref.data(), ref.size());
 }
 
 /// Creates a StringRef out of its MLIR C API equivalent.
 inline llvm::StringRef unwrap(MlirStringRef ref) {
   return llvm::StringRef(ref.data, ref.length);
 }
 
 inline MlirLogicalResult wrap(mlir::LogicalResult res) {
   if (mlir::succeeded(res))
     return mlirLogicalResultSuccess();
   return mlirLogicalResultFailure();
 }
 
 inline mlir::LogicalResult unwrap(MlirLogicalResult res) {
   return mlir::success(mlirLogicalResultIsSuccess(res));
 }
 
+DEFINE_C_API_PTR_METHODS(MlirLlvmThreadPool, llvm::ThreadPool)
 DEFINE_C_API_METHODS(MlirTypeID, mlir::TypeID)
 DEFINE_C_API_PTR_METHODS(MlirTypeIDAllocator, mlir::TypeIDAllocator)
 
 #endif // MLIR_CAPI_SUPPORT_H

mlir/lib/CAPI/IR/IR.cpp

 // Context API.
 //===----------------------------------------------------------------------===//
 
 MlirContext mlirContextCreate() {
   auto *context = new MLIRContext;
   return wrap(context);
 }
 
+static inline MLIRContext::Threading toThreadingEnum(bool threadingEnabled) {
+  return threadingEnabled ? MLIRContext::Threading::ENABLED
+                          : MLIRContext::Threading::DISABLED;
+}
+
+MlirContext mlirContextCreateWithThreading(bool threadingEnabled) {
+  auto *context = new MLIRContext(toThreadingEnum(threadingEnabled));
+  return wrap(context);
+}
+
+MlirContext mlirContextCreateWithRegistry(MlirDialectRegistry registry,
+                                          bool threadingEnabled) {
+  auto *context =
+      new MLIRContext(*unwrap(registry), toThreadingEnum(threadingEnabled));
+  return wrap(context);
+}
+
 bool mlirContextEqual(MlirContext ctx1, MlirContext ctx2) {
   return unwrap(ctx1) == unwrap(ctx2);
 }
 
 void mlirContextDestroy(MlirContext context) { delete unwrap(context); }
 
 void mlirContextSetAllowUnregisteredDialects(MlirContext context, bool allow) {
   unwrap(context)->allowUnregisteredDialects(allow);

 void mlirContextEnableMultithreading(MlirContext context, bool enable) {
   return unwrap(context)->enableMultithreading(enable);
 }
 
 void mlirContextLoadAllAvailableDialects(MlirContext context) {
   unwrap(context)->loadAllAvailableDialects();
 }
 
+void mlirContextSetThreadPool(MlirContext context,
+                              MlirLlvmThreadPool threadPool) {
+  unwrap(context)->setThreadPool(*unwrap(threadPool));
+}
+
 //===----------------------------------------------------------------------===//
 // Dialect API.
 //===----------------------------------------------------------------------===//
 
 MlirContext mlirDialectGetContext(MlirDialect dialect) {
   return wrap(unwrap(dialect)->getContext());
 }
 

mlir/lib/CAPI/IR/Support.cpp

 //===- Support.cpp - Helpers for C interface to MLIR API ------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/CAPI/Support.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/ThreadPool.h"
 
 #include <cstring>
 
 MlirStringRef mlirStringRefCreateFromCString(const char *str) {
   return mlirStringRefCreate(str, strlen(str));
 }
 
 bool mlirStringRefEqual(MlirStringRef string, MlirStringRef other) {
   return llvm::StringRef(string.data, string.length) ==
          llvm::StringRef(other.data, other.length);
 }
 
 //===----------------------------------------------------------------------===//
+// LLVM ThreadPool API.
+//===----------------------------------------------------------------------===//
+MlirLlvmThreadPool mlirLlvmThreadPoolCreate() {
+  return wrap(new llvm::ThreadPool());
+}
+
+void mlirLlvmThreadPoolDestroy(MlirLlvmThreadPool threadPool) {
+  delete unwrap(threadPool);
+}
+
+//===----------------------------------------------------------------------===//
 // TypeID API.
 //===----------------------------------------------------------------------===//
 MlirTypeID mlirTypeIDCreate(const void *ptr) {
   assert(reinterpret_cast<uintptr_t>(ptr) % 8 == 0 &&
          "ptr must be 8 byte aligned");
   // This is essentially a no-op that returns back `ptr`, but by going through
   // the `TypeID` functions we can get compiler errors in case the `TypeID`
   // api/representation changes

mlir/test/CAPI/ir.c

int testDialectRegistry(void) {
   }
 
   mlirContextDestroy(ctx);
   mlirDialectRegistryDestroy(registry);
 
   return 0;
 }
 
+void testExplicitThreadPools(void) {
+  MlirLlvmThreadPool threadPool = mlirLlvmThreadPoolCreate();
+  MlirDialectRegistry registry = mlirDialectRegistryCreate();
+  mlirRegisterAllDialects(registry);
+  MlirContext context =
+      mlirContextCreateWithRegistry(registry, /*threadingEnabled=*/false);
+  mlirContextSetThreadPool(context, threadPool);
+  mlirContextDestroy(context);
+  mlirDialectRegistryDestroy(registry);
+  mlirLlvmThreadPoolDestroy(threadPool);
+}
+
 void testDiagnostics(void) {
   MlirContext ctx = mlirContextCreate();
   MlirDiagnosticHandlerID id = mlirContextAttachDiagnosticHandler(
       ctx, errorHandler, (void *)42, deleteUserData);
   fprintf(stderr, "@test_diagnostics\n");
   MlirLocation unknownLoc = mlirLocationUnknownGet(ctx);
   mlirEmitError(unknownLoc, "test diagnostics");
   MlirAttribute unknownAttr = mlirLocationGetAttribute(unknownLoc);
if (testTypeID(ctx))
     return 13;
   if (testSymbolTable(ctx))
     return 14;
   if (testDialectRegistry())
     return 15;
 
   mlirContextDestroy(ctx);
 
+  testExplicitThreadPools();
   testDiagnostics();
   return 0;
 }