[MLIR][Python] enable precise registration

Author: makslevental

mlir/include/mlir-c/Bindings/Python/Interop.h

@@ -84,6 +84,8 @@
 #define MLIR_PYTHON_CAPSULE_VALUE MAKE_MLIR_PYTHON_QUALNAME("ir.Value._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_TYPEID                                             \
   MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_DIALECT_HANDLE                                     \
+  MAKE_MLIR_PYTHON_QUALNAME("ir.DialectHandle._CAPIPtr")
 
 /** Attribute on MLIR Python objects that expose their C-API pointer.
  * This will be a type-specific capsule created as per one of the helpers
@@ -457,6 +459,13 @@ static inline MlirValue mlirPythonCapsuleToValue(PyObject *capsule) {
   return value;
 }
 
+static inline MlirDialectHandle
+mlirPythonCapsuleToDialectHandle(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_DIALECT_HANDLE);
+  MlirDialectHandle handle = {ptr};
+  return handle;
+}
+
 #ifdef __cplusplus
 }
 #endif

mlir/include/mlir-c/Dialect/Builtin.h

@@ -0,0 +1,33 @@
+//===-- mlir-c/Dialect/Builtin.h - C API for Builtin dialect ------*- C -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header declares the C interface for registering and accessing the
+// Builtin dialect. A dialect should be registered with a context to make it
+// available to users of the context. These users must load the dialect
+// before using any of its attributes, operations or types. Parser and pass
+// manager can load registered dialects automatically.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_C_DIALECT_BUILTIN_H
+#define MLIR_C_DIALECT_BUILTIN_H
+
+#include "mlir-c/IR.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Builtin, builtin);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif // MLIR_C_DIALECT_BUILTIN_H

mlir/include/mlir-c/IR.h

@@ -66,6 +66,7 @@ DEFINE_C_API_STRUCT(MlirLocation, const void);
 DEFINE_C_API_STRUCT(MlirModule, const void);
 DEFINE_C_API_STRUCT(MlirType, const void);
 DEFINE_C_API_STRUCT(MlirValue, const void);
+DEFINE_C_API_STRUCT(MlirDialectHandle, const void);
 
 #undef DEFINE_C_API_STRUCT
 
@@ -207,11 +208,6 @@ MLIR_CAPI_EXPORTED MlirStringRef mlirDialectGetNamespace(MlirDialect dialect);
 // registration schemes.
 //===----------------------------------------------------------------------===//
 
-struct MlirDialectHandle {
-  const void *ptr;
-};
-typedef struct MlirDialectHandle MlirDialectHandle;
-
 #define MLIR_DECLARE_CAPI_DIALECT_REGISTRATION(Name, Namespace)                \
   MLIR_CAPI_EXPORTED MlirDialectHandle mlirGetDialectHandle__##Namespace##__(  \
       void)
@@ -233,6 +229,11 @@ MLIR_CAPI_EXPORTED void mlirDialectHandleRegisterDialect(MlirDialectHandle,
 MLIR_CAPI_EXPORTED MlirDialect mlirDialectHandleLoadDialect(MlirDialectHandle,
                                                             MlirContext);
 
+/// Checks if the dialect handle is null.
+static inline bool mlirDialectHandleIsNull(MlirDialectHandle handle) {
+  return !handle.ptr;
+}
+
 //===----------------------------------------------------------------------===//
 // DialectRegistry API.
 //===----------------------------------------------------------------------===//
@@ -249,6 +250,13 @@ static inline bool mlirDialectRegistryIsNull(MlirDialectRegistry registry) {
 MLIR_CAPI_EXPORTED void
 mlirDialectRegistryDestroy(MlirDialectRegistry registry);
 
+MLIR_CAPI_EXPORTED int64_t
+mlirDialectRegistryGetNumDialectNames(MlirDialectRegistry registry);
+
+MLIR_CAPI_EXPORTED void
+mlirDialectRegistryGetDialectNames(MlirDialectRegistry registry,
+                                   MlirStringRef *dialectNames);
+
 //===----------------------------------------------------------------------===//
 // Location API.
 //===----------------------------------------------------------------------===//

mlir/lib/Bindings/Python/IRCore.cpp

@@ -2897,6 +2897,14 @@ maybeGetTracebackLocation(const std::optional<PyLocation> &location) {
 // Populates the core exports of the 'ir' submodule.
 //------------------------------------------------------------------------------
 
+MlirDialectHandle createMlirDialectHandleFromCapsule(nb::object capsule) {
+  MlirDialectHandle rawRegistry =
+      mlirPythonCapsuleToDialectHandle(capsule.ptr());
+  if (mlirDialectHandleIsNull(rawRegistry))
+    throw nb::python_error();
+  return rawRegistry;
+}
+
 void mlir::python::populateIRCore(nb::module_ &m) {
   // disable leak warnings which tend to be false positives.
   nb::set_leak_warnings(false);
@@ -3126,14 +3134,38 @@ void mlir::python::populateIRCore(nb::module_ &m) {
           },
           nb::sig("def __repr__(self) -> str"));
 
+  //----------------------------------------------------------------------------
+  // Mapping of MlirDialectHandle
+  //----------------------------------------------------------------------------
+
+  nb::class_<MlirDialectHandle>(m, "DialectHandle")
+      .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,
+                  &createMlirDialectHandleFromCapsule);
+
   //----------------------------------------------------------------------------
   // Mapping of PyDialectRegistry
   //----------------------------------------------------------------------------
   nb::class_<PyDialectRegistry>(m, "DialectRegistry")
       .def_prop_ro(MLIR_PYTHON_CAPI_PTR_ATTR, &PyDialectRegistry::getCapsule)
       .def_static(MLIR_PYTHON_CAPI_FACTORY_ATTR,
                   &PyDialectRegistry::createFromCapsule)
-      .def(nb::init<>());
+      .def(nb::init<>())
+      .def("insert_dialect",
+           [](PyDialectRegistry &self, MlirDialectHandle handle) {
+             mlirDialectHandleInsertDialect(handle, self.get());
+           })
+      .def("insert_dialect",
+           [](PyDialectRegistry &self, intptr_t ptr) {
+             mlirDialectHandleInsertDialect(
+                 {reinterpret_cast<const void *>(ptr)}, self.get());
+           })
+      .def_prop_ro("dialect_names", [](PyDialectRegistry &self) {
+        int64_t numDialectNames =
+            mlirDialectRegistryGetNumDialectNames(self.get());
+        std::vector<MlirStringRef> dialectNames(numDialectNames);
+        mlirDialectRegistryGetDialectNames(self.get(), dialectNames.data());
+        return dialectNames;
+      });
 
   //----------------------------------------------------------------------------
   // Mapping of Location

mlir/lib/CAPI/Dialect/Builtin.cpp

@@ -0,0 +1,13 @@
+//===- Builtin.cpp - C Interface for Builtin dialect ----------------------===//
+//
+// 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-c/Dialect/Builtin.h"
+#include "mlir/CAPI/Registration.h"
+#include "mlir/IR/BuiltinDialect.h"
+
+MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Builtin, builtin, mlir::BuiltinDialect)

mlir/lib/CAPI/Dialect/CMakeLists.txt

@@ -16,6 +16,14 @@ add_mlir_upstream_c_api_library(MLIRCAPIArith
   MLIRArithDialect
 )
 
+add_mlir_upstream_c_api_library(MLIRCAPIBuiltin
+  Builtin.cpp
+
+  PARTIAL_SOURCES_INTENDED
+  LINK_LIBS PUBLIC
+  MLIRCAPIIR
+)
+
 add_mlir_upstream_c_api_library(MLIRCAPIAsync
   Async.cpp
   AsyncPasses.cpp

mlir/lib/CAPI/IR/IR.cpp

@@ -150,6 +150,18 @@ void mlirDialectRegistryDestroy(MlirDialectRegistry registry) {
   delete unwrap(registry);
 }
 
+int64_t mlirDialectRegistryGetNumDialectNames(MlirDialectRegistry registry) {
+  auto dialectNames = unwrap(registry)->getDialectNames();
+  return std::distance(dialectNames.begin(), dialectNames.end());
+}
+
+void mlirDialectRegistryGetDialectNames(MlirDialectRegistry registry,
+                                        MlirStringRef *dialectNames) {
+  for (auto [i, location] :
+       llvm::enumerate(unwrap(registry)->getDialectNames()))
+    dialectNames[i] = wrap(location);
+}
+
 //===----------------------------------------------------------------------===//
 // AsmState API.
 //===----------------------------------------------------------------------===//

mlir/python/CMakeLists.txt

@@ -528,31 +528,32 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     MLIRCAPIDebug
     MLIRCAPIIR
     MLIRCAPIInterfaces
+    MLIRCAPITransforms
+    MLIRCAPIBuiltin
 
     # Dialects
     MLIRCAPIFunc
 )
 
-# This extension exposes an API to register all dialects, extensions, and passes
-# packaged in upstream MLIR and it is used for the upstream "mlir" Python
-# package. Downstreams will likely want to provide their own and not depend
-# on this one, since it links in the world.
-# Note that this is not added to any top-level source target for transitive
-# inclusion: It must be included explicitly by downstreams if desired. Note that
-# this has a very large impact on what gets built/packaged.
-declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
-  MODULE_NAME _mlirRegisterEverything
-  ROOT_DIR "${PYTHON_SOURCE_DIR}"
-  PYTHON_BINDINGS_LIBRARY nanobind
-  SOURCES
-    RegisterEverything.cpp
-  PRIVATE_LINK_LIBS
-    LLVMSupport
-  EMBED_CAPI_LINK_LIBS
-    MLIRCAPIConversion
-    MLIRCAPITransforms
-    MLIRCAPIRegisterEverything
-)
+## This extension exposes an API to register all dialects, extensions, and passes
+## packaged in upstream MLIR and it is used for the upstream "mlir" Python
+## package. Downstreams will likely want to provide their own and not depend
+## on this one, since it links in the world.
+## Note that this is not added to any top-level source target for transitive
+## inclusion: It must be included explicitly by downstreams if desired. Note that
+## this has a very large impact on what gets built/packaged.
+#declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
+#  MODULE_NAME _mlirRegisterEverything
+#  ROOT_DIR "${PYTHON_SOURCE_DIR}"
+#  PYTHON_BINDINGS_LIBRARY nanobind
+#  SOURCES
+#    RegisterEverything.cpp
+#  PRIVATE_LINK_LIBS
+#    LLVMSupport
+#    MLIRCAPIConversion
+#    MLIRCAPITransforms
+#    MLIRCAPIRegisterEverything
+#)
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
   MODULE_NAME _mlirDialectsLinalg
@@ -871,7 +872,6 @@ add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME}
     MLIRPythonCAPI.HeaderSources
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.RegisterEverything
     ${_ADDL_TEST_SOURCES}
 )
 
@@ -979,7 +979,6 @@ add_mlir_python_modules(MLIRPythonModules
   INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
   DECLARED_SOURCES
     MLIRPythonSources
-    MLIRPythonExtension.RegisterEverything
     MLIRPythonExtension.Core.type_stub_gen
     MLIRPythonCAPICTypesBinding
     ${_ADDL_TEST_SOURCES}

mlir/python/mlir/_mlir_libs/_capi.py.in

@@ -5,4 +5,12 @@
 import ctypes
 from pathlib import Path
 
-_capi = ctypes.CDLL(str(Path(__file__).parent / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"))
+_capi = ctypes.CDLL(str(Path(__file__).parent / "@CMAKE_SHARED_LIBRARY_PREFIX@@MLIR_PYTHON_CAPI_DYLIB_NAME@@CMAKE_SHARED_LIBRARY_SUFFIX@"))
+
+PyCapsule_New = ctypes.pythonapi.PyCapsule_New
+PyCapsule_New.restype = ctypes.py_object
+PyCapsule_New.argtypes = ctypes.c_void_p, ctypes.c_char_p, ctypes.c_void_p
+
+MLIR_PYTHON_CAPSULE_DIALECT_HANDLE = (
+    "@[email protected]._CAPIPtr"
+).encode()

mlir/test/python/dialects/irdl.py

@@ -2,6 +2,7 @@
 
 from mlir.ir import *
 from mlir.dialects.irdl import *
+import mlir.dialects.arith
 import sys