diff --git a/mlir/cmake/modules/AddMLIRPython.cmake b/mlir/cmake/modules/AddMLIRPython.cmake --- a/mlir/cmake/modules/AddMLIRPython.cmake +++ b/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 @@ 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 @@ # 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) diff --git a/mlir/python/CMakeLists.txt b/mlir/python/CMakeLists.txt --- a/mlir/python/CMakeLists.txt +++ b/mlir/python/CMakeLists.txt @@ -46,6 +46,12 @@ 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 @@ DECLARED_SOURCES MLIRPythonSources MLIRPythonExtension.AllPassesRegistration + MLIRPythonCAPIHeaderSources COMMON_CAPI_LINK_LIBS MLIRPythonCAPI ) diff --git a/mlir/python/mlir/_mlir_libs/__init__.py b/mlir/python/mlir/_mlir_libs/__init__.py --- a/mlir/python/mlir/_mlir_libs/__init__.py +++ b/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 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")] diff --git a/mlir/test/python/develoment_files.py b/mlir/test/python/develoment_files.py new file mode 100644 --- /dev/null +++ b/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}"