[mlir][capi] Add DialectRegistry to MLIR C-API

[trilorez]

Summary

Exposes mlir::DialectRegistry to the C API as MlirDialectRegistry along with
helper functions. A hook has been added to MlirDialectHandle that inserts
the dialect into a registry.

[GeorgeLyon]

mlir/include/mlir/CAPI/Registration.h

31

We can probably remove this and implement the C function in terms of dialectregistry

[stellaraccident]

mlir/include/mlir/CAPI/Registration.h

31

Agreed. I believe this is used in a number of downstreams. Can you document the recommended change in the commit description and note that we will remove this soon. Add a todo here.

[GeorgeLyon]

mlir/include/mlir/CAPI/Registration.h

31

I’m pro deprecation, but I was merely referring to the function pointer. We can keep the C API and just call the registry version from the implementation. So we could remove this and add a note.

[trilorez]

Added TODO regarding removing MlirContextRegisterDialectHook and updated
commit message.

[trilorez]

mlir/include/mlir/CAPI/Registration.h

31

Oh, I interpreted this as people were directly using the function pointer. I can get rid of it if that isn't the case.

[trilorez]

Remove MlirContextRegisterDialectHook and implement
mlirDialectHandleRegisterDialect using mlirDialectHandleInsertDialect

[trilorez]

mlir/include/mlir/CAPI/Registration.h

31

Updated to remove the function pointer and implement mlirDialectHandleRegisterDialect using mlirDialectHandleInsertDialect.

[trilorez]

Fix MLIR_DEFINE_CAPI_DIALECT_REGISTRATION

[mehdi_amini]

mlir/include/mlir-c/Registration.h

54

Do we need to deprecate this API? Seems like it can stay as a convenient wrapper?

[trilorez]

mlir/include/mlir-c/Registration.h

54

That's true, I can just leave the declaration as is for now

[trilorez]

Remove deprecation TODO

mlir/include/mlir-c/IR.h

 #define DEFINE_C_API_STRUCT(name, storage)                                     \
   struct name {                                                                \
     storage *ptr;                                                              \
   };                                                                           \
   typedef struct name name
 
 DEFINE_C_API_STRUCT(MlirContext, void);
 DEFINE_C_API_STRUCT(MlirDialect, void);
+DEFINE_C_API_STRUCT(MlirDialectRegistry, void);
 DEFINE_C_API_STRUCT(MlirOperation, void);
 DEFINE_C_API_STRUCT(MlirOpPrintingFlags, void);
 DEFINE_C_API_STRUCT(MlirBlock, void);
 DEFINE_C_API_STRUCT(MlirRegion, void);
 DEFINE_C_API_STRUCT(MlirSymbolTable, void);
 
 DEFINE_C_API_STRUCT(MlirAttribute, const void);
 DEFINE_C_API_STRUCT(MlirIdentifier, const void);

 MLIR_CAPI_EXPORTED bool
 mlirContextGetAllowUnregisteredDialects(MlirContext context);
 
 /// Returns the number of dialects registered with the given context. A
 /// registered dialect will be loaded if needed by the parser.
 MLIR_CAPI_EXPORTED intptr_t
 mlirContextGetNumRegisteredDialects(MlirContext context);
 
+/// Append the contents of the given dialect registry to the registry associated
+/// with the context.
+MLIR_CAPI_EXPORTED void
+mlirContextAppendDialectRegistry(MlirContext ctx, MlirDialectRegistry registry);
+
 /// Returns the number of dialects loaded by the context.
 
 MLIR_CAPI_EXPORTED intptr_t
 mlirContextGetNumLoadedDialects(MlirContext context);
 
 /// Gets the dialect instance owned by the given context using the dialect
 /// namespace to identify it, loads (i.e., constructs the instance of) the
 /// dialect if necessary. If the dialect is not registered with the context,

 /// from different contexts will not compare equal.
 MLIR_CAPI_EXPORTED bool mlirDialectEqual(MlirDialect dialect1,
                                          MlirDialect dialect2);
 
 /// Returns the namespace of the given dialect.
 MLIR_CAPI_EXPORTED MlirStringRef mlirDialectGetNamespace(MlirDialect dialect);
 
 //===----------------------------------------------------------------------===//
+// DialectRegistry API.
+//===----------------------------------------------------------------------===//
+
+/// Creates a dialect registry and transfers its ownership to the caller.
+MLIR_CAPI_EXPORTED MlirDialectRegistry mlirDialectRegistryCreate();
+
+/// Checks if the dialect registry is null.
+static inline bool mlirDialectRegistryIsNull(MlirDialectRegistry registry) {
+  return !registry.ptr;
+}
+
+/// Takes a dialect registry owned by the caller and destroys it.
+MLIR_CAPI_EXPORTED void
+mlirDialectRegistryDestroy(MlirDialectRegistry registry);
+
+//===----------------------------------------------------------------------===//
 // Location API.
 //===----------------------------------------------------------------------===//
 
 /// Creates an File/Line/Column location owned by the given context.
 MLIR_CAPI_EXPORTED MlirLocation mlirLocationFileLineColGet(
     MlirContext context, MlirStringRef filename, unsigned line, unsigned col);
 
 /// Creates a call site location with a callee and a caller.

mlir/include/mlir-c/Registration.h

 
 #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);
 
+/// Inserts the dialect associated with the provided dialect handle into the
+/// provided dialect registry
+MLIR_CAPI_EXPORTED void mlirDialectHandleInsertDialect(MlirDialectHandle,
+                                                       MlirDialectRegistry);
+
 /// 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/include/mlir/CAPI/IR.h

 
 #include "mlir/CAPI/Wrap.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
 
 DEFINE_C_API_PTR_METHODS(MlirContext, mlir::MLIRContext)
 DEFINE_C_API_PTR_METHODS(MlirDialect, mlir::Dialect)
+DEFINE_C_API_PTR_METHODS(MlirDialectRegistry, mlir::DialectRegistry)
 DEFINE_C_API_PTR_METHODS(MlirOperation, mlir::Operation)
 DEFINE_C_API_PTR_METHODS(MlirBlock, mlir::Block)
 DEFINE_C_API_PTR_METHODS(MlirOpPrintingFlags, mlir::OpPrintingFlags)
 DEFINE_C_API_PTR_METHODS(MlirRegion, mlir::Region)
 DEFINE_C_API_PTR_METHODS(MlirSymbolTable, mlir::SymbolTable)
 
 DEFINE_C_API_METHODS(MlirAttribute, mlir::Attribute)
 DEFINE_C_API_METHODS(MlirIdentifier, mlir::StringAttr)
 DEFINE_C_API_METHODS(MlirLocation, mlir::Location)
 DEFINE_C_API_METHODS(MlirModule, mlir::ModuleOp)
 DEFINE_C_API_METHODS(MlirType, mlir::Type)
 DEFINE_C_API_METHODS(MlirTypeID, mlir::TypeID)
 DEFINE_C_API_METHODS(MlirValue, mlir::Value)
 
 #endif // MLIR_CAPI_IR_H

mlir/include/mlir/CAPI/Registration.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 void (*MlirDialectRegistryInsertDialectHook)(
+    MlirDialectRegistry registry);
 typedef MlirDialect (*MlirContextLoadDialectHook)(MlirContext context);
 typedef MlirStringRef (*MlirDialectGetNamespaceHook)();
 
 /// Structure of dialect registration hooks.
 struct MlirDialectRegistrationHooks {
-  MlirContextRegisterDialectHook registerHook;
+  MlirDialectRegistryInsertDialectHook insertHook;
   MlirContextLoadDialectHook loadHook;
   MlirDialectGetNamespaceHook getNamespaceHook;
 };
 typedef struct MlirDialectRegistrationHooks MlirDialectRegistrationHooks;
 
 #define MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Name, Namespace, ClassName)      \
-  static void mlirContextRegister##Name##Dialect(MlirContext context) {        \
-    mlir::DialectRegistry registry;                                            \
-    registry.insert<ClassName>();                                              \
-    unwrap(context)->appendDialectRegistry(registry);                          \
+  static void mlirDialectRegistryInsert##Name##Dialect(                        \
+      MlirDialectRegistry registry) {                                          \
+    unwrap(registry)->insert<ClassName>();                                     \
   }                                                                            \
   static MlirDialect mlirContextLoad##Name##Dialect(MlirContext context) {     \
     return wrap(unwrap(context)->getOrLoadDialect<ClassName>());               \
   }                                                                            \
   static MlirStringRef mlir##Name##DialectGetNamespace() {                     \
     return wrap(ClassName::getDialectNamespace());                             \
   }                                                                            \
   MlirDialectHandle mlirGetDialectHandle__##Namespace##__() {                  \
     static MlirDialectRegistrationHooks hooks = {                              \
-        mlirContextRegister##Name##Dialect, mlirContextLoad##Name##Dialect,    \
-        mlir##Name##DialectGetNamespace};                                      \
+        mlirDialectRegistryInsert##Name##Dialect,                              \
+        mlirContextLoad##Name##Dialect, mlir##Name##DialectGetNamespace};      \
     return MlirDialectHandle{&hooks};                                          \
   }
 
 #endif // MLIR_CAPI_REGISTRATION_H

mlir/lib/CAPI/IR/DialectHandle.cpp

 unwrap(MlirDialectHandle handle) {
   return (const MlirDialectRegistrationHooks *)handle.ptr;
 }
 
 MlirStringRef mlirDialectHandleGetNamespace(MlirDialectHandle handle) {
   return unwrap(handle)->getNamespaceHook();
 }
 
+void mlirDialectHandleInsertDialect(MlirDialectHandle handle,
+                                    MlirDialectRegistry registry) {
+  unwrap(handle)->insertHook(registry);
+}
+
 void mlirDialectHandleRegisterDialect(MlirDialectHandle handle,
                                       MlirContext ctx) {
-  unwrap(handle)->registerHook(ctx);
+  mlir::DialectRegistry registry;
+  mlirDialectHandleInsertDialect(handle, wrap(&registry));
+  unwrap(ctx)->appendDialectRegistry(registry);
 }
 
 MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle handle,
                                          MlirContext ctx) {
   return unwrap(handle)->loadHook(ctx);
 }

mlir/lib/CAPI/IR/IR.cpp

 
 bool mlirContextGetAllowUnregisteredDialects(MlirContext context) {
   return unwrap(context)->allowsUnregisteredDialects();
 }
 intptr_t mlirContextGetNumRegisteredDialects(MlirContext context) {
   return static_cast<intptr_t>(unwrap(context)->getAvailableDialects().size());
 }
 
+void mlirContextAppendDialectRegistry(MlirContext ctx,
+                                      MlirDialectRegistry registry) {
+  unwrap(ctx)->appendDialectRegistry(*unwrap(registry));
+}
+
 // TODO: expose a cheaper way than constructing + sorting a vector only to take
 // its size.
 intptr_t mlirContextGetNumLoadedDialects(MlirContext context) {
   return static_cast<intptr_t>(unwrap(context)->getLoadedDialects().size());
 }
 
 MlirDialect mlirContextGetOrLoadDialect(MlirContext context,
                                         MlirStringRef name) {
bool mlirDialectEqual(MlirDialect dialect1, MlirDialect dialect2) {
   return unwrap(dialect1) == unwrap(dialect2);
 }
 
 MlirStringRef mlirDialectGetNamespace(MlirDialect dialect) {
   return wrap(unwrap(dialect)->getNamespace());
 }
 
 //===----------------------------------------------------------------------===//
+// DialectRegistry API.
+//===----------------------------------------------------------------------===//
+
+MlirDialectRegistry mlirDialectRegistryCreate() {
+  return wrap(new DialectRegistry());
+}
+
+void mlirDialectRegistryDestroy(MlirDialectRegistry registry) {
+  delete unwrap(registry);
+}
+
+//===----------------------------------------------------------------------===//
 // Printing flags API.
 //===----------------------------------------------------------------------===//
 
 MlirOpPrintingFlags mlirOpPrintingFlagsCreate() {
   return wrap(new OpPrintingFlags());
 }
 
 void mlirOpPrintingFlagsDestroy(MlirOpPrintingFlags flags) {

mlir/test/CAPI/ir.c

int testSymbolTable(MlirContext ctx) {
 
   mlirSymbolTableDestroy(symbolTable);
   mlirModuleDestroy(module);
   mlirModuleDestroy(otherModule);
 
   return 0;
 }
 
+int testDialectRegistry() {
+  fprintf(stderr, "@testDialectRegistry\n");
+
+  MlirDialectRegistry registry = mlirDialectRegistryCreate();
+  if (mlirDialectRegistryIsNull(registry)) {
+    fprintf(stderr, "ERROR: Expected registry to be present\n");
+    return 1;
+  }
+
+  MlirDialectHandle stdHandle = mlirGetDialectHandle__std__();
+  mlirDialectHandleInsertDialect(stdHandle, registry);
+
+  MlirContext ctx = mlirContextCreate();
+  if (mlirContextGetNumRegisteredDialects(ctx) != 0) {
+    fprintf(stderr,
+            "ERROR: Expected no dialects to be registered to new context\n");
+  }
+
+  mlirContextAppendDialectRegistry(ctx, registry);
+  if (mlirContextGetNumRegisteredDialects(ctx) != 1) {
+    fprintf(stderr, "ERROR: Expected the dialect in the registry to be "
+                    "registered to the context\n");
+  }
+
+  mlirContextDestroy(ctx);
+  mlirDialectRegistryDestroy(registry);
+
+  return 0;
+}
+
 void testDiagnostics() {
   MlirContext ctx = mlirContextCreate();
   MlirDiagnosticHandlerID id = mlirContextAttachDiagnosticHandler(
       ctx, errorHandler, (void *)42, deleteUserData);
   fprintf(stderr, "@test_diagnostics\n");
   MlirLocation unknownLoc = mlirLocationUnknownGet(ctx);
   mlirEmitError(unknownLoc, "test diagnostics");
   MlirLocation fileLineColLoc = mlirLocationFileLineColGet(
int main() {
   if (testOperands())
     return 11;
   if (testClone())
     return 12;
   if (testTypeID(ctx))
     return 13;
   if (testSymbolTable(ctx))
     return 14;
+  if (testDialectRegistry())
+    return 15;
 
   mlirContextDestroy(ctx);
 
   testDiagnostics();
   return 0;
 }