[MLIR][Python] enable precise registration

Author: makslevental

mlir/cmake/modules/AddMLIRPython.cmake

@@ -23,11 +23,14 @@
 #     grouping. Source groupings form a DAG.
 #   SOURCES: List of specific source files relative to ROOT_DIR to include.
 #   SOURCES_GLOB: List of glob patterns relative to ROOT_DIR to include.
+#   EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
+#     on. These will be collected for all extensions and put into an
+#     aggregate dylib that is linked against.
 function(declare_mlir_python_sources name)
   cmake_parse_arguments(ARG
     ""
     "ROOT_DIR;ADD_TO_PARENT"
-    "SOURCES;SOURCES_GLOB"
+    "SOURCES;SOURCES_GLOB;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
 
   if(NOT ARG_ROOT_DIR)
@@ -53,9 +56,10 @@ function(declare_mlir_python_sources name)
   set_target_properties(${name} PROPERTIES
     # Yes: Leading-lowercase property names are load bearing and the recommended
     # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
-    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_DEPENDS"
+    EXPORT_PROPERTIES "mlir_python_SOURCES_TYPE;mlir_python_DEPENDS;mlir_python_EMBED_CAPI_LINK_LIBS"
     mlir_python_SOURCES_TYPE pure
     mlir_python_DEPENDS ""
+    mlir_python_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
   )
 
   # Use the interface include directories and sources on the target to carry the
@@ -374,6 +378,9 @@ endfunction()
 #     This file is where the *EnumAttrs are defined, not where the *Enums are defined.
 #     **WARNING**: This arg will shortly be removed when the just-below TODO is satisfied. Use at your
 #     risk.
+#   EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
+#     on. These will be collected for all extensions and put into an
+#     aggregate dylib that is linked against.
 #
 # TODO: Right now `TD_FILE` can't be the actual dialect tablegen file, since we
 #       use its path to determine where to place the generated python file. If
@@ -383,7 +390,7 @@ function(declare_mlir_dialect_python_bindings)
   cmake_parse_arguments(ARG
     "GEN_ENUM_BINDINGS"
     "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
-    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE"
+    "SOURCES;SOURCES_GLOB;DEPENDS;GEN_ENUM_BINDINGS_TD_FILE;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
   # Sources.
   set(_dialect_target "${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}")
@@ -424,6 +431,7 @@ function(declare_mlir_dialect_python_bindings)
       ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
       ADD_TO_PARENT "${_dialect_target}"
       SOURCES ${_sources}
+      EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
     )
   endif()
 endfunction()

mlir/python/CMakeLists.txt

@@ -326,7 +326,10 @@ declare_mlir_dialect_python_bindings(
   SOURCES
     dialects/arith.py
   DIALECT_NAME arith
-  GEN_ENUM_BINDINGS)
+  GEN_ENUM_BINDINGS
+  EMBED_CAPI_LINK_LIBS
+    MLIRCAPIArith
+  )
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -846,8 +849,20 @@ endif()
 # once ready.
 ################################################################################
 
+set(MLIR_PYTHON_CAPI_DYLIB_NAME MLIRPythonCAPI)
+configure_file(
+  "${CMAKE_CURRENT_LIST_DIR}/mlir/_mlir_libs/_capi.py.in"
+  "${CMAKE_CURRENT_BINARY_DIR}/_mlir_libs/_capi.py"
+  @ONLY
+)
+declare_mlir_python_sources(
+  MLIRPythonCAPICTypesBinding
+  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
+  SOURCES _mlir_libs/_capi.py
+)
+
 set(MLIRPythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}")
-add_mlir_python_common_capi_library(MLIRPythonCAPI
+add_mlir_python_common_capi_library(${MLIR_PYTHON_CAPI_DYLIB_NAME}
   INSTALL_COMPONENT MLIRPythonModules
   INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
   OUTPUT_DIRECTORY "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs"
@@ -966,6 +981,7 @@ add_mlir_python_modules(MLIRPythonModules
     MLIRPythonSources
     MLIRPythonExtension.RegisterEverything
     MLIRPythonExtension.Core.type_stub_gen
+    MLIRPythonCAPICTypesBinding
     ${_ADDL_TEST_SOURCES}
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI

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

@@ -0,0 +1,8 @@
+#  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
+
+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@"))

mlir/test/python/ir/capi.py

@@ -0,0 +1,6 @@
+# RUN: %PYTHON %s | FileCheck %s
+
+from mlir._mlir_libs._capi import _capi
+
+print("success")
+# CHECK: success