[MLIR][python] Export CAPI headers.

[stellaraccident]

Summary

• 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).

[mikeurbach]

Makes sense to me.

mlir/cmake/modules/AddMLIRPython.cmake

 #     For relocatable sources, this will point deeper into the directory that
 #     can be relocated. For generated sources, can be relative to
 #     CMAKE_CURRENT_BINARY_DIR. Generated and non generated sources cannot be
 #     mixed.
 #   ADD_TO_PARENT: Adds this source grouping to a previously declared source
 #     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})
 
   if(NOT ARG_ROOT_DIR)
     set(ARG_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
   endif()
 
   # Process the glob.
function(declare_mlir_python_sources name)
   # We create a custom target to carry properties and dependencies for
   # generated sources.
   add_custom_target(${name})
   set(_file_depends "${ARG_SOURCES}")
   list(TRANSFORM _file_depends PREPEND "${ARG_ROOT_DIR}/")
   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 ""
   )
 
   # Add to parent.
   if(ARG_ADD_TO_PARENT)
     set_property(TARGET ${ARG_ADD_TO_PARENT} APPEND PROPERTY PYTHON_DEPENDS ${name})
cmake_parse_arguments(ARG
     ${ARGN})
   # Helper to process an individual target.
   function(_process_target modules_target sources_target)
     get_target_property(_source_type ${sources_target} PYTHON_SOURCES_TYPE)
     if(_source_type STREQUAL "pure")
       # 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)
 
         file(MAKE_DIRECTORY "${_dest_dir}")
         add_custom_command(
           TARGET ${modules_target} PRE_BUILD
           COMMENT "Copying python source ${_src_path} -> ${_dest_path}"

mlir/python/CMakeLists.txt

declare_mlir_python_sources(MLIRPythonSources.Passes
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   ADD_TO_PARENT MLIRPythonSources
   SOURCES_GLOB
     all_passes_registration/*.py
     conversions/*.py
     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
 ################################################################################
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   TD_FILE dialects/AsyncOps.td

 ################################################################################
 
 add_mlir_python_modules(MLIRPythonModules
   ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
   INSTALL_PREFIX "python_packages/mlir_core/mlir"
   DECLARED_SOURCES
     MLIRPythonSources
     MLIRPythonExtension.AllPassesRegistration
+    MLIRPythonCAPIHeaderSources
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
   )
 
 
 add_mlir_python_modules(MLIRPythonTestModules
   ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_test/mlir"
   INSTALL_PREFIX "python_packages/mlir_test/mlir"
   DECLARED_SOURCES
     MLIRPythonTestSources
   )

mlir/python/mlir/_mlir_libs/__init__.py

 # Licensed 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
 
+from typing import Sequence
+
 import importlib
 import os
 
 __all__ = [
   "load_extension",
   "preload_dependency",
 ]
 
 _this_dir = os.path.dirname(__file__)
 
 def load_extension(name):
   return importlib.import_module(f".{name}", __package__)
 
 
 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

+# 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}"