Skip to content

Commit

Permalink
guard python passes behind LLVM_ENABLE_RTTI and AIE_ENABLE_PYTHON_PASS
Browse files Browse the repository at this point in the history
  • Loading branch information
makslevental committed Nov 18, 2023
1 parent 533114c commit f59cfc5
Show file tree
Hide file tree
Showing 19 changed files with 482 additions and 274 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/buildAndTest.yml
Original file line number Diff line number Diff line change
Expand Up @@ -157,7 +157,7 @@ jobs:
max-size: 1G

- name: Build and install LLVM
run: utils/build-llvm.sh
run: LLVM_ENABLE_RTTI=ON utils/build-llvm.sh

# Build the repo test target in debug mode to build and test.
- name: Build and test (Assert)
Expand Down
34 changes: 23 additions & 11 deletions CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -125,18 +125,28 @@ set(LLVM_LIT_ARGS

list(APPEND CMAKE_MODULE_PATH "${MLIR_CMAKE_DIR}")
list(APPEND CMAKE_MODULE_PATH "${LLVM_CMAKE_DIR}")

include(TableGen)
include(AddLLVM)
include(AddMLIR)
include(HandleLLVMOptions)
include(ExternalProject)
include(CMakeDependentOption)

include_directories(${LLVM_INCLUDE_DIRS})
include_directories(${MLIR_INCLUDE_DIRS})
include_directories(${PROJECT_SOURCE_DIR}/include)
include_directories(${PROJECT_BINARY_DIR}/include)
add_definitions(${LLVM_DEFINITIONS})


cmake_dependent_option(AIE_ENABLE_BINDINGS_PYTHON
"Enables building of Python bindings." ON "MLIR_ENABLE_BINDINGS_PYTHON" OFF)

cmake_dependent_option(AIE_ENABLE_PYTHON_PASSES
"Enables building of passes that connect to python." ON "AIE_ENABLE_BINDINGS_PYTHON;LLVM_ENABLE_RTTI" OFF)


# Silence a false positive GCC -Wunused-but-set-parameter warning in constexpr
# cases, by marking SelectedCase as used. See
# https://gcc.gnu.org/bugzilla/show_bug.cgi?id=85827 for details. The issue is
Expand All @@ -159,18 +169,17 @@ if(AIE_INCLUDE_DOCS)
add_dependencies(docs mlir-doc)
endif()

# setup python
include(MLIRDetectPythonEnv)
mlir_configure_python_dev_packages()
if(AIE_ENABLE_BINDINGS_PYTHON)
include(MLIRDetectPythonEnv)
mlir_configure_python_dev_packages()

# python install directory
if(NOT AIE_PYTHON_PACKAGES_DIR)
set(AIE_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python")
endif()
if(NOT AIE_PYTHON_PACKAGES_DIR)
set(AIE_PYTHON_PACKAGES_DIR "${CMAKE_CURRENT_BINARY_DIR}/python")
endif()

# python install directory
if(NOT AIE_PYTHON_INSTALL_DIR)
set(AIE_PYTHON_INSTALL_DIR "python")
if(NOT AIE_PYTHON_INSTALL_DIR)
set(AIE_PYTHON_INSTALL_DIR "python")
endif()
endif()

add_subdirectory(include)
Expand All @@ -179,7 +188,10 @@ add_subdirectory(runtime_lib)
add_subdirectory(aie_runtime_lib)
add_subdirectory(tools)
add_subdirectory(data)
add_subdirectory(python)

if(AIE_ENABLE_BINDINGS_PYTHON)
add_subdirectory(python)
endif()

if(NOT LLVM_INSTALL_TOOLCHAIN_ONLY)
install(
Expand Down
1 change: 1 addition & 0 deletions cmake/modules/AIEConfig.cmake.in
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ set(AIE_CMAKE_DIR "@AIE_CONFIG_CMAKE_DIR@")
set(AIE_BINARY_DIR "@AIE_CONFIG_BINARY_DIR@")
set(AIE_TOOLS_BINARY_DIR "@AIE_CONFIG_TOOLS_BINARY_DIR@")
set(AIE_INCLUDE_DIRS "@AIE_CONFIG_INCLUDE_DIRS@")
set(LLVM_ENABLE_RTTI @LLVM_ENABLE_RTTI@)

# Provide AIE configurations for users to inherit from.
set(AIE_RUNTIME_TARGETS "@AIE_RUNTIME_TARGETS@")
Expand Down
6 changes: 0 additions & 6 deletions python/AIEMLIRModule.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,6 @@
//
//===----------------------------------------------------------------------===//

#include "PythonPass.h"

#include "aie-c/Dialects.h"
#include "aie-c/Registration.h"

Expand Down Expand Up @@ -46,10 +44,6 @@ PYBIND11_MODULE(_aie, m) {
},
py::arg("context"), py::arg("load") = true);

m.def("register_python_pass_demo_pass", [](py::function func) {
registerPythonPassDemoPassWithFunc(std::move(func));
});

// AIE types bindings
mlir_type_subclass(m, "ObjectFifoType", aieTypeIsObjectFifoType)
.def_classmethod(
Expand Down
236 changes: 151 additions & 85 deletions python/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,8 @@ declare_mlir_python_sources(AIEPythonSources
ROOT_DIR "${AIE_PYTHON_ROOT_DIR}"
)

declare_mlir_python_sources(AIEPythonExtensions)

declare_mlir_python_sources(AIEPythonSources.Dialects
ADD_TO_PARENT AIEPythonSources
)
Expand Down Expand Up @@ -47,61 +49,151 @@ declare_mlir_dialect_python_bindings(
# Extensions
################################################################################

get_property(MLIR_CONVERSION_LIBS GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
get_property(MLIR_DIALECT_LIBS GLOBAL PROPERTY MLIR_DIALECT_LIBS)
get_property(MLIR_EXTENSION_LIBS GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
get_property(MLIR_TRANSLATION_LIBS GLOBAL PROPERTY MLIR_TRANSLATION_LIBS)
get_property(MLIR_UPSTREAM_CAPI_LIBS GLOBAL PROPERTY MLIR_UPSTREAM_CAPI_LIBS)

get_target_property(RegisterEverythingSources MLIRPythonExtension.RegisterEverything INTERFACE_SOURCES)

_flatten_mlir_python_targets(mlir_python_sources MLIRPythonSources)
set(_mlir_python_cpp_sources)
foreach(s ${mlir_python_sources})
get_target_property(_sources ${s} INTERFACE_SOURCES)
list(FILTER _sources INCLUDE REGEX "(\.cpp|\.h)$")
if(_sources)
list(APPEND _mlir_python_cpp_sources ${_sources})
endif()
endforeach()
if (AIE_ENABLE_PYTHON_PASSES)
# What's going on here? This branch builds a single shared library Python C extension module
# with absolutely everything, including the kitchen sink. I.e, this "omnibus" shared library
# will have absolutely all the MLIR code (including C APIs), all the upstream C extension modules,
# our main C extension module (AIEMLIRModule.cpp) and the PythonPass.cpp C extension module.
#
# Why do this? This is a way to call arbitrary MLIR code from the python bindings. Such as, for example,
# PythonPass.cpp does.
#
# Note, this depends on MLIR being compiled with RTTI (i.e., LLVM_ENABLE_RTTI=ON) and the
# guard includes this.

get_property(MLIR_CONVERSION_LIBS GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
get_property(MLIR_DIALECT_LIBS GLOBAL PROPERTY MLIR_DIALECT_LIBS)
get_property(MLIR_EXTENSION_LIBS GLOBAL PROPERTY MLIR_EXTENSION_LIBS)
get_property(MLIR_TRANSLATION_LIBS GLOBAL PROPERTY MLIR_TRANSLATION_LIBS)
get_property(MLIR_UPSTREAM_CAPI_LIBS GLOBAL PROPERTY MLIR_UPSTREAM_CAPI_LIBS)

# Dig out all of the upstream cpp source files associated with the C extension modules.
_flatten_mlir_python_targets(mlir_python_sources MLIRPythonSources)
set(_mlir_python_cpp_sources)
foreach(s ${mlir_python_sources})
get_target_property(_sources ${s} INTERFACE_SOURCES)
list(FILTER _sources INCLUDE REGEX "(\.cpp|\.h)$")
if(_sources)
list(APPEND _mlir_python_cpp_sources ${_sources})
endif()
endforeach()

# Missed one...
get_target_property(RegisterEverythingSources MLIRPythonExtension.RegisterEverything INTERFACE_SOURCES)

declare_mlir_python_extension(AIEPythonExtensions.MLIR
MODULE_NAME _aie
ADD_TO_PARENT AIEPythonExtensions
ROOT_DIR "/"

SOURCES

${CMAKE_CURRENT_SOURCE_DIR}/AIEMLIRModule.cpp
${CMAKE_CURRENT_SOURCE_DIR}/PythonPass.cpp

${_mlir_python_cpp_sources}
${RegisterEverythingSources}

PRIVATE_LINK_LIBS
LLVMSupport
AIECAPI
${MLIR_CONVERSION_LIBS}
${MLIR_DIALECT_LIBS}
${MLIR_EXTENSION_LIBS}
${MLIR_TRANSLATION_LIBS}
${MLIR_UPSTREAM_CAPI_LIBS}
)

list(SORT _mlir_python_cpp_sources)
# Kill all of the dependencies (mlir_python_DEPENDS) on cpp sources in the dialect build targets
# so the C extension modules don't get built but the Python sources do get installed.
_flatten_mlir_python_targets(mlir_python_sources_deps MLIRPythonSources)
list(FILTER mlir_python_sources_deps INCLUDE REGEX "^MLIRPythonSources.Dialects")
foreach(t ${mlir_python_sources_deps})
set_property(TARGET ${t} PROPERTY mlir_python_DEPENDS "")
endforeach()
set_property(TARGET MLIRPythonSources.ExecutionEngine PROPERTY mlir_python_DEPENDS "")

add_mlir_python_modules(AIEPythonModules
ROOT_PREFIX "${AIE_PYTHON_PACKAGES_DIR}/aie"
INSTALL_PREFIX "python/aie"
DECLARED_SOURCES
${mlir_python_sources_deps}
MLIRPythonSources.Core.Python
MLIRPythonSources.ExecutionEngine
AIEPythonExtensions.MLIR
AIEPythonSources
)

declare_mlir_python_extension(AIEPythonExtensions
MODULE_NAME _aie
ROOT_DIR "/"
set(_other_extensions
_aie_python_passes
_mlir
_mlirAsyncPasses
_mlirDialectsLinalg
_mlirDialectsPDL
_mlirDialectsQuant
_mlirDialectsSparseTensor
_mlirDialectsTransform
_mlirExecutionEngine
_mlirGPUPasses
_mlirLinalgPasses
_mlirPythonTest
_mlirRegisterEverything
_mlirSparseTensorPasses
)

file(MAKE_DIRECTORY "${CMAKE_BINARY_DIR}/python/aie/_mlir_libs")

# This is the trick that makes multiple Python C extension modules packaged into one shared library work;
# Python importlib machinery looks for PyInit__your_module based on the filename _your_module.so.
# Thus, if you have a symlink _your_module.so -> BIGOMNIBUS.so, it will look for PyInit__your_module in
# BIGOMNIBUS.so.
foreach(ext ${_other_extensions})
# Build time/directory creation.
add_custom_target("symlink_${ext}"
COMMAND ${CMAKE_COMMAND} -E create_symlink
"_aie${PYTHON_MODULE_EXTENSION}"
"${ext}${PYTHON_MODULE_EXTENSION}"
DEPENDS AIEPythonExtensions
WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/python/aie/_mlir_libs"
)
add_dependencies(AIEPythonModules "symlink_${ext}")
# Install time/directory creation.
install(CODE "execute_process(
COMMAND ${CMAKE_COMMAND} -E create_symlink
_aie${PYTHON_MODULE_EXTENSION}
${ext}${PYTHON_MODULE_EXTENSION}
WORKING_DIRECTORY ${CMAKE_INSTALL_PREFIX}/python/aie/_mlir_libs
)"
)
endforeach()

else ()

# Standard Python C extension bindings.

declare_mlir_python_extension(AIEPythonExtensions.MLIR
MODULE_NAME _aie
ADD_TO_PARENT AIEPythonExtensions
ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR}
SOURCES
${CMAKE_CURRENT_SOURCE_DIR}/AIEMLIRModule.cpp
${CMAKE_CURRENT_SOURCE_DIR}/PythonPass.cpp

${_mlir_python_cpp_sources}
${RegisterEverythingSources}

AIEMLIRModule.cpp
EMBED_CAPI_LINK_LIBS
AIECAPI
PRIVATE_LINK_LIBS
LLVMSupport
AIECAPI
${MLIR_CONVERSION_LIBS}
${MLIR_DIALECT_LIBS}
${MLIR_EXTENSION_LIBS}
${MLIR_TRANSLATION_LIBS}
${MLIR_UPSTREAM_CAPI_LIBS}
LLVMSupport
)

_flatten_mlir_python_targets(mlir_python_sources_deps MLIRPythonSources)
list(FILTER mlir_python_sources_deps INCLUDE REGEX "^MLIRPythonSources.Dialects")
foreach(t ${mlir_python_sources_deps})
set_property(TARGET ${t} PROPERTY mlir_python_DEPENDS "")
endforeach()
set_property(TARGET MLIRPythonSources.ExecutionEngine PROPERTY mlir_python_DEPENDS "")

add_mlir_python_modules(AIEMLIRPythonModules
ROOT_PREFIX "${AIE_PYTHON_PACKAGES_DIR}/aie"
INSTALL_PREFIX "python/aie"
add_mlir_python_common_capi_library(AIEAggregateCAPI
INSTALL_COMPONENT AIEPythonModules
INSTALL_DESTINATION python/aie/_mlir_libs
OUTPUT_DIRECTORY "${AIE_PYTHON_PACKAGES_DIR}/aie/_mlir_libs"
RELATIVE_INSTALL_ROOT "../../../.."
DECLARED_SOURCES
${mlir_python_sources_deps}
MLIRPythonSources.Core.Python
MLIRPythonSources.ExecutionEngine
MLIRPythonSources
MLIRPythonExtension.Core
MLIRPythonExtension.RegisterEverything
MLIRPythonExtension.ExecutionEngine
AIEPythonSources
AIEPythonExtensions
)

Expand All @@ -110,47 +202,19 @@ add_mlir_python_modules(AIEPythonModules
INSTALL_PREFIX "python/aie"
DECLARED_SOURCES
AIEPythonSources
MLIRPythonSources
MLIRPythonExtension.Core
MLIRPythonExtension.RegisterEverything
MLIRPythonExtension.ExecutionEngine
AIEPythonExtensions
COMMON_CAPI_LINK_LIBS
AIEAggregateCAPI
AIECAPI
)

add_subdirectory(compiler)
add_dependencies(AIEPythonModules AIECompilerPythonModules)

set(_other_extensions
_mlir
_mlirAsyncPasses
_mlirDialectsLinalg
_mlirDialectsPDL
_mlirDialectsQuant
_mlirDialectsSparseTensor
_mlirDialectsTransform
_mlirExecutionEngine
_mlirGPUPasses
_mlirLinalgPasses
_mlirPythonTest
_mlirRegisterEverything
_mlirSparseTensorPasses
)

file(MAKE_DIRECTORY "${CMAKE_BINARY_DIR}/python/aie/_mlir_libs")

foreach(ext ${_other_extensions})
add_custom_target("symlink_${ext}"
COMMAND ${CMAKE_COMMAND} -E create_symlink
"_aie${PYTHON_MODULE_EXTENSION}"
"${ext}${PYTHON_MODULE_EXTENSION}"
DEPENDS AIEPythonExtensions
WORKING_DIRECTORY "${CMAKE_BINARY_DIR}/python/aie/_mlir_libs"
)
add_dependencies(AIEPythonModules "symlink_${ext}")
install(CODE "execute_process(
COMMAND ${CMAKE_COMMAND} -E create_symlink
_aie${PYTHON_MODULE_EXTENSION}
${ext}${PYTHON_MODULE_EXTENSION}
WORKING_DIRECTORY ${CMAKE_INSTALL_PREFIX}/python/aie/_mlir_libs
)"
)
endforeach()
endif ()

# Copy the runtime libs into the _mlir_libs directory for convenience.
set(_runtimes
mlir_async_runtime
mlir_c_runner_utils
Expand All @@ -177,3 +241,5 @@ install(IMPORTED_RUNTIME_ARTIFACTS
LIBRARY DESTINATION "python/aie/_mlir_libs"
)

add_subdirectory(compiler)
add_dependencies(AIEPythonModules AIECompilerPythonModules)
Loading

0 comments on commit f59cfc5

Please sign in to comment.