[MLIR] Replace dialect registration hooks with dialect handle

[GeorgeLyon]

Summary

Replace MlirDialectRegistrationHooks with MlirDialectHandle, which under-the-hood is an opaque pointer to MlirDialectRegistrationHooks. Then we expose the functionality previously directly on MlirDialectRegistrationHooks, as functions which take the opaque MlirDialectHandle struct. This makes the actual structure of the registration hooks an implementation detail, and happens to avoid this issue: https://llvm.discourse.group/t/strange-swift-issues-with-dialect-registration-hooks/2759/3

[GeorgeLyon]

mlir/include/mlir-c/Registration.h

41

We can also just make this a symbol instead of a function.

[stellaraccident]

mlir/lib/CAPI/IR/DialectHandle.cpp

12

Can we match naming used elsewhere, even if the mechanism is different: "unwrapHooks". Have a look at Support.h too: I suspect that you can just do unwrap(handle) if you've defined the hooks struct properly.

29

Newline please.

[GeorgeLyon]

Address feedback

[GeorgeLyon]

mlir/lib/CAPI/IR/DialectHandle.cpp

12

I'm not an expert in C++, but I'm not sure this would work as MlirDialectRegistrationHooks is a C type. I have changed the name instead.

[GeorgeLyon]

mlir/lib/CAPI/IR/DialectHandle.cpp

12

Nevermind, I thought you meant something different. Just changing the name to unwrap works.

[GeorgeLyon]

Slight tewak

mlir/include/mlir-c/Registration.h

 
 //===----------------------------------------------------------------------===//
 // Dialect registration declarations.
 // Registration entry-points for each dialect are declared using the common
 // MLIR_DECLARE_DIALECT_REGISTRATION_CAPI macro, which takes the dialect
 // API name (i.e. "Standard", "Tensor", "Linalg") and namespace (i.e. "std",
 // "tensor", "linalg"). The following declarations are produced:
 //
-//   /// Registers the dialect with the given context. This allows the
-//   /// dialect to be loaded dynamically if needed when parsing. */
-//   void mlirContextRegister{NAME}Dialect(MlirContext);
-//
-//   /// Loads the dialect into the given context. The dialect does _not_
-//   /// have to be registered in advance.
-//   MlirDialect mlirContextLoad{NAME}Dialect(MlirContext context);
-//
-//   /// Returns the namespace of the Standard dialect, suitable for loading it.
-//   MlirStringRef mlir{NAME}DialectGetNamespace();
-//
 //   /// Gets the above hook methods in struct form for a dialect by namespace.
 //   /// This is intended to facilitate dynamic lookup and registration of
 //   /// dialects via a plugin facility based on shared library symbol lookup.
-//   const MlirDialectRegistrationHooks *mlirGetDialectHooks__{NAMESPACE}__();
+//   const MlirDialectHandle *mlirGetDialectHandle__{NAMESPACE}__();
 //
 // This is done via a common macro to facilitate future expansion to
 // registration schemes.
 //===----------------------------------------------------------------------===//
 
-#define MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Name, Namespace)                \
-  MLIR_CAPI_EXPORTED void mlirContextRegister##Name##Dialect(                  \
-      MlirContext context);                                                    \
-  MLIR_CAPI_EXPORTED MlirDialect mlirContextLoad##Name##Dialect(               \
-      MlirContext context);                                                    \
-  MLIR_CAPI_EXPORTED MlirStringRef mlir##Name##DialectGetNamespace();          \
-  MLIR_CAPI_EXPORTED const MlirDialectRegistrationHooks                        \
-      *mlirGetDialectHooks__##Namespace##__()
-
-/// Hooks for dynamic discovery of dialects.
-typedef void (*MlirContextRegisterDialectHook)(MlirContext context);
-typedef MlirDialect (*MlirContextLoadDialectHook)(MlirContext context);
-typedef MlirStringRef (*MlirDialectGetNamespaceHook)();
-
-/// Structure of dialect registration hooks.
-struct MlirDialectRegistrationHooks {
-  MlirContextRegisterDialectHook registerHook;
-  MlirContextLoadDialectHook loadHook;
-  MlirDialectGetNamespaceHook getNamespaceHook;
+struct MlirDialectHandle {
+  const void *ptr;
 };
-typedef struct MlirDialectRegistrationHooks MlirDialectRegistrationHooks;
+typedef struct MlirDialectHandle MlirDialectHandle;
+
+#define MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Name, Namespace)                \
+  MLIR_CAPI_EXPORTED MlirDialectHandle mlirGetDialectHandle__##Namespace##__()
+
+/// Returns the namespace associated with the provided dialect handle.
+MLIR_CAPI_EXPORTED
+MlirStringRef mlirDialectHandleGetNamespace(MlirDialectHandle);
+
+/// Registers the dialect associated with the provided dialect handle.
+MLIR_CAPI_EXPORTED void mlirDialectHandleRegisterDialect(MlirDialectHandle,
+                                                         MlirContext);
+
+/// Loads the dialect associated with the provided dialect handle.
+MLIR_CAPI_EXPORTED MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle,
+                                                            MlirContext);
 
 /// Registers all dialects known to core MLIR with the provided Context.
 /// This is needed before creating IR for these Dialects.
 /// TODO: Remove this function once the real registration API is finished.
 MLIR_CAPI_EXPORTED void mlirRegisterAllDialects(MlirContext context);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_REGISTRATION_H

mlir/include/mlir/CAPI/Registration.h

 #include "mlir/CAPI/Support.h"
 
 //===----------------------------------------------------------------------===//
 // Corrolary to MLIR_DECLARE_CAPI_DIALECT_REGISTRATION that defines an impl.
 // Takes the same name passed to the above and the fully qualified class name
 // of the dialect class.
 //===----------------------------------------------------------------------===//
 
+/// Hooks for dynamic discovery of dialects.
+typedef void (*MlirContextRegisterDialectHook)(MlirContext context);
+typedef MlirDialect (*MlirContextLoadDialectHook)(MlirContext context);
+typedef MlirStringRef (*MlirDialectGetNamespaceHook)();
+
+/// Structure of dialect registration hooks.
+struct MlirDialectRegistrationHooks {
+  MlirContextRegisterDialectHook registerHook;
+  MlirContextLoadDialectHook loadHook;
+  MlirDialectGetNamespaceHook getNamespaceHook;
+};
+typedef struct MlirDialectRegistrationHooks MlirDialectRegistrationHooks;
+
 #define MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Name, Namespace, ClassName)      \
-  void mlirContextRegister##Name##Dialect(MlirContext context) {               \
+  static void mlirContextRegister##Name##Dialect(MlirContext context) {        \
     unwrap(context)->getDialectRegistry().insert<ClassName>();                 \
   }                                                                            \
-  MlirDialect mlirContextLoad##Name##Dialect(MlirContext context) {            \
+  static MlirDialect mlirContextLoad##Name##Dialect(MlirContext context) {     \
     return wrap(unwrap(context)->getOrLoadDialect<ClassName>());               \
   }                                                                            \
-  MlirStringRef mlir##Name##DialectGetNamespace() {                            \
+  static MlirStringRef mlir##Name##DialectGetNamespace() {                     \
     return wrap(ClassName::getDialectNamespace());                             \
   }                                                                            \
-  const MlirDialectRegistrationHooks *mlirGetDialectHooks__##Namespace##__() { \
+  MlirDialectHandle mlirGetDialectHandle__##Namespace##__() {                  \
     static MlirDialectRegistrationHooks hooks = {                              \
         mlirContextRegister##Name##Dialect, mlirContextLoad##Name##Dialect,    \
         mlir##Name##DialectGetNamespace};                                      \
-    return &hooks;                                                             \
+    return MlirDialectHandle{&hooks};                                          \
   }
 
 #endif // MLIR_CAPI_REGISTRATION_H

mlir/lib/CAPI/IR/CMakeLists.txt

 # Main API shared library.
 add_mlir_public_c_api_library(MLIRCAPIIR
   AffineExpr.cpp
   AffineMap.cpp
   BuiltinAttributes.cpp
   BuiltinTypes.cpp
   Diagnostics.cpp
+  DialectHandle.cpp
   IntegerSet.cpp
   IR.cpp
   Pass.cpp
   Support.cpp
 
   LINK_LIBS PUBLIC
   MLIRIR
   MLIRParser
   MLIRSupport
   MLIRPass
 )

mlir/lib/CAPI/IR/DialectHandle.cpp

+//===- DialectHandle.cpp - C Interface for MLIR Dialect Operations -------===//
+//
+// 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/Registration.h"
+
+static inline const MlirDialectRegistrationHooks *
+unwrap(MlirDialectHandle handle) {
+  return (const MlirDialectRegistrationHooks *)handle.ptr;
+}
+
+MlirStringRef mlirDialectHandleGetNamespace(MlirDialectHandle handle) {
+  return unwrap(handle)->getNamespaceHook();
+}
+
+void mlirDialectHandleRegisterDialect(MlirDialectHandle handle,
+                                      MlirContext ctx) {
+  unwrap(handle)->registerHook(ctx);
+}
+
+MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle handle,
+                                         MlirContext ctx) {
+  return unwrap(handle)->loadHook(ctx);
+}

mlir/test/CAPI/ir.c

 }
 
 int registerOnlyStd() {
   MlirContext ctx = mlirContextCreate();
   // The built-in dialect is always loaded.
   if (mlirContextGetNumLoadedDialects(ctx) != 1)
     return 1;
 
-  MlirDialect std =
-      mlirContextGetOrLoadDialect(ctx, mlirStandardDialectGetNamespace());
+  MlirDialectHandle stdHandle = mlirGetDialectHandle__std__();
+
+  MlirDialect std = mlirContextGetOrLoadDialect(
+      ctx, mlirDialectHandleGetNamespace(stdHandle));
   if (!mlirDialectIsNull(std))
     return 2;
 
-  mlirContextRegisterStandardDialect(ctx);
+  mlirDialectHandleRegisterDialect(stdHandle, ctx);
 
-  std = mlirContextGetOrLoadDialect(ctx, mlirStandardDialectGetNamespace());
+  std = mlirContextGetOrLoadDialect(ctx,
+                                    mlirDialectHandleGetNamespace(stdHandle));
   if (mlirDialectIsNull(std))
     return 3;
 
-  MlirDialect alsoStd = mlirContextLoadStandardDialect(ctx);
+  MlirDialect alsoStd = mlirDialectHandleLoadDialect(stdHandle, ctx);
   if (!mlirDialectEqual(std, alsoStd))
     return 4;
 
   MlirStringRef stdNs = mlirDialectGetNamespace(std);
-  MlirStringRef alsoStdNs = mlirStandardDialectGetNamespace();
+  MlirStringRef alsoStdNs = mlirDialectHandleGetNamespace(stdHandle);
   if (stdNs.length != alsoStdNs.length ||
       strncmp(stdNs.data, alsoStdNs.data, stdNs.length))
     return 5;
 
   fprintf(stderr, "@registration\n");
   // CHECK-LABEL: @registration
 
   return 0;