torch-mlir/python/CMakeLists.txt

153 lines
5.2 KiB
CMake

include(AddMLIRPython)
include(MLIRDetectPythonEnv)
################################################################################
# Resources that must be packaged into the python tree
################################################################################
file(MAKE_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/npcomp/compiler/backend/refjit_resources")
add_custom_target(NPCOMPPythonResources ALL)
add_custom_command(
TARGET NPCOMPPythonResources
COMMAND ${CMAKE_COMMAND} -E copy
# TODO: Make the runtime library work for windows.
# TODO: Use $<TARGET-FILE:> for this.
${CMAKE_BINARY_DIR}/lib/libNPCOMPCompilerRuntimeShlib${CMAKE_SHARED_LIBRARY_SUFFIX}
${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core/npcomp/compiler/generic/backend/libNPCOMPCompilerRuntimeShlib${CMAKE_SHARED_LIBRARY_SUFFIX}
)
add_dependencies(NPCOMPPythonResources
NPCOMPCompilerRuntimeShlib
)
################################################################################
# Declare sources
################################################################################
declare_mlir_python_sources(NPCOMPPythonSources)
declare_mlir_python_sources(NPCOMPPythonSources.Core
ADD_TO_PARENT NPCOMPPythonSources
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
SOURCES
npcomp/__init__.py
npcomp/build.py
npcomp/decorators.py
npcomp/exporter.py
npcomp/smoketest.py
npcomp/types.py
npcomp/dialects/_ods_common.py
SOURCES_GLOB
npcomp/compiler/*.py
npcomp/frontends/*.py
npcomp/torch/*.py
npcomp/tracing/*.py
npcomp/utils/*.py
)
declare_mlir_python_sources(NPCOMPPythonSources.Dialects
ADD_TO_PARENT NPCOMPPythonSources
)
declare_mlir_python_sources(NPCOMPPythonExtensions)
declare_mlir_python_sources(NPCOMPPythonCAPIHeaderSources
ROOT_DIR "${MLIR_NPCOMP_SOURCE_DIR}/include"
SOURCES_GLOB "npcomp-c/*.h"
DEST_PREFIX "_mlir_libs/include"
)
################################################################################
# Extensions
################################################################################
set(_addl_extension_sources)
if(NPCOMP_ENABLE_REFJIT)
list(APPEND _addl_extension_sources "${CMAKE_CURRENT_SOURCE_DIR}/RefJITBackend.cpp")
endif()
declare_mlir_python_extension(NPCOMPPythonExtensions.Core
MODULE_NAME _npcomp
ADD_TO_PARENT NPCOMPPythonExtensions
SOURCES
${CMAKE_CURRENT_SOURCE_DIR}/NpcompModule.cpp
${_addl_extension_sources}
EMBED_CAPI_LINK_LIBS
NPCOMPCAPI
PRIVATE_LINK_LIBS
LLVMSupport
)
################################################################################
# Declare dialects
################################################################################
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT NPCOMPPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
TD_FILE npcomp/dialects/BasicpyBind.td
SOURCES npcomp/dialects/basicpy.py
DIALECT_NAME basicpy)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT NPCOMPPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
TD_FILE npcomp/dialects/NumpyBind.td
SOURCES npcomp/dialects/numpy.py
DIALECT_NAME numpy)
declare_mlir_dialect_python_bindings(
ADD_TO_PARENT NPCOMPPythonSources.Dialects
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}"
TD_FILE npcomp/dialects/TorchBind.td
SOURCES npcomp/dialects/torch.py
DIALECT_NAME torch)
################################################################################
# Build composite binaries
################################################################################
# Bundle our own, self-contained CAPI library with all of our deps.
add_mlir_python_common_capi_library(NPCOMPPythonCAPI
INSTALL_COMPONENT NPCOMPPythonModules
INSTALL_DESTINATION python_packages/npcomp_core/mlir/_mlir_libs
# NOTE: When the MLIR API is relocated under npcomp, this would change to
# .../npcomp/_mlir_libs
OUTPUT_DIRECTORY "${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core/mlir/_mlir_libs"
RELATIVE_INSTALL_ROOT "../../../.."
DECLARED_SOURCES
# TODO: This can be chopped down significantly for size.
MLIRPythonSources
MLIRPythonExtension.AllPassesRegistration
NPCOMPPythonSources
NPCOMPPythonExtensions
)
# Bundle the MLIR python sources into our package.
# The MLIR API is position independent, so we explicitly output it to the mlir/
# folder as a temporary measure. It will eventually migrate under the npcomp/
# folder and be accessible under the unified "import npcomp..." namespace.
add_mlir_python_modules(NPCOMPMLIRPythonModules
ROOT_PREFIX "${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core/mlir"
INSTALL_PREFIX "python_packages/npcomp_core/mlir"
DECLARED_SOURCES
MLIRPythonSources
MLIRPythonExtension.AllPassesRegistration
# We need the npcomp extensions co-located with the MLIR extensions. When
# the namespace is unified, this moves to the below.
MLIRPythonCAPIHeaderSources
NPCOMPPythonExtensions
NPCOMPPythonCAPIHeaderSources
COMMON_CAPI_LINK_LIBS
NPCOMPPythonCAPI
)
# Bundle the NPCOMP python sources into our package.
add_mlir_python_modules(NPCOMPPythonModules
ROOT_PREFIX "${MLIR_NPCOMP_PYTHON_PACKAGES_DIR}/npcomp_core"
INSTALL_PREFIX "python_packages/npcomp_core"
DECLARED_SOURCES
NPCOMPPythonSources
COMMON_CAPI_LINK_LIBS
NPCOMPPythonCAPI
)