[MLIR][python] Export CAPI headers.

* Adds source targets (not included in the full set that downstreams use by default) to bundle mlir-c/ headers into the mlir/_mlir_libs/include directory.
* Adds a minimal entry point to get include and library directories.
* Used by npcomp to export a full CAPI (which is then used by the Torch extension to link npcomp).

Reviewed By: mikeurbach

Differential Revision: https://reviews.llvm.org/D107090

Author: stellaraccident

mlir/cmake/modules/AddMLIRPython.cmake

@@ -23,10 +23,12 @@
 #     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.
+#   DEST_PREFIX: Destination prefix to prepend to files in the python
+#     package directory namespace.
 function(declare_mlir_python_sources name)
   cmake_parse_arguments(ARG
     ""
-    "ROOT_DIR;ADD_TO_PARENT"
+    "ROOT_DIR;ADD_TO_PARENT;DEST_PREFIX"
     "SOURCES;SOURCES_GLOB"
     ${ARGN})
 
@@ -54,6 +56,7 @@ function(declare_mlir_python_sources name)
   set_target_properties(${name} PROPERTIES
     PYTHON_SOURCES_TYPE pure
     PYTHON_ROOT_DIR "${ARG_ROOT_DIR}"
+    PYTHON_DEST_PREFIX "${ARG_DEST_PREFIX}"
     PYTHON_SOURCES "${ARG_SOURCES}"
     PYTHON_FILE_DEPENDS "${_file_depends}"
     PYTHON_DEPENDS ""
@@ -130,9 +133,14 @@ function(add_mlir_python_modules name)
       # Pure python sources to link into the tree.
       get_target_property(_python_root_dir ${sources_target} PYTHON_ROOT_DIR)
       get_target_property(_python_sources ${sources_target} PYTHON_SOURCES)
+      get_target_property(_specified_dest_prefix ${sources_target} PYTHON_DEST_PREFIX)
+      set(_dest_prefix "${ARG_ROOT_PREFIX}")
+      if(_specified_dest_prefix)
+        set(_dest_prefix "${_dest_prefix}/${_specified_dest_prefix}")
+      endif()
       foreach(_source_relative_path ${_python_sources})
         set(_src_path "${_python_root_dir}/${_source_relative_path}")
-        set(_dest_path "${ARG_ROOT_PREFIX}/${_source_relative_path}")
+        set(_dest_path "${_dest_prefix}/${_source_relative_path}")
 
         get_filename_component(_dest_dir "${_dest_path}" DIRECTORY)
         get_filename_component(_install_path "${ARG_INSTALL_PREFIX}/${_source_relative_path}" DIRECTORY)

mlir/python/CMakeLists.txt

@@ -46,6 +46,12 @@ declare_mlir_python_sources(MLIRPythonSources.Passes
     transforms/*.py
 )
 
+declare_mlir_python_sources(MLIRPythonCAPIHeaderSources
+  ROOT_DIR "${MLIR_SOURCE_DIR}/include"
+  SOURCES_GLOB "mlir-c/*.h"
+  DEST_PREFIX "_mlir_libs/include"
+)
+
 ################################################################################
 # Dialect bindings
 ################################################################################
@@ -304,6 +310,7 @@ add_mlir_python_modules(MLIRPythonModules
   DECLARED_SOURCES
     MLIRPythonSources
     MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonCAPIHeaderSources
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
   )

mlir/python/mlir/_mlir_libs/__init__.py

@@ -2,6 +2,8 @@
 # See https://llvm.org/LICENSE.txt for license information.
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
+from typing import Sequence
+
 import importlib
 import os
 
@@ -19,3 +21,21 @@ def load_extension(name):
 def preload_dependency(public_name):
   # TODO: Implement this hook to pre-load DLLs with ctypes on Windows.
   pass
+
+
+def get_lib_dirs() -> Sequence[str]:
+  """Gets the lib directory for linking to shared libraries.
+
+  On some platforms, the package may need to be built specially to export
+  development libraries.
+  """
+  return [_this_dir]
+
+
+def get_include_dirs() -> Sequence[str]:
+  """Gets the include directory for compiling against exported C libraries.
+
+  Depending on how the package was build, development C libraries may or may
+  not be present.
+  """
+  return [os.path.join(_this_dir, "include")]

mlir/test/python/develoment_files.py

@@ -0,0 +1,18 @@
+# RUN: %PYTHON %s 2>&1
+
+import os
+
+from mlir._mlir_libs import get_include_dirs, get_lib_dirs
+
+
+header_file = os.path.join(get_include_dirs()[0], "mlir-c", "IR.h")
+assert os.path.isfile(header_file), f"Header does not exist: {header_file}"
+
+# Since actual library names are platform specific, just scan the directory
+# for a filename that contains the library name.
+expected_lib_name = "MLIRPythonCAPI"
+all_libs = os.listdir(get_lib_dirs()[0])
+found_lib = False
+for file_name in all_libs:
+  if expected_lib_name in file_name: found_lib = True
+assert found_lib, f"Did not find '{expected_lib_name}' lib in {all_libs}"