1################################################################################
2# Python modules
3# MLIR's Python modules are both directly used by the core project and are
4# available for use and embedding into external projects (in their own
5# namespace and with their own deps). In order to facilitate this, python
6# artifacts are split between declarations, which make a subset of
7# things available to be built and "add", which in line with the normal LLVM
8# nomenclature, adds libraries.
9################################################################################
10
11# Function: declare_mlir_python_sources
12# Declares pure python sources as part of a named grouping that can be built
13# later.
14# Arguments:
15#   ROOT_DIR: Directory where the python namespace begins (defaults to
16#     CMAKE_CURRENT_SOURCE_DIR). For non-relocatable sources, this will
17#     typically just be the root of the python source tree (current directory).
18#     For relocatable sources, this will point deeper into the directory that
19#     can be relocated. For generated sources, can be relative to
20#     CMAKE_CURRENT_BINARY_DIR. Generated and non generated sources cannot be
21#     mixed.
22#   ADD_TO_PARENT: Adds this source grouping to a previously declared source
23#     grouping. Source groupings form a DAG.
24#   SOURCES: List of specific source files relative to ROOT_DIR to include.
25#   SOURCES_GLOB: List of glob patterns relative to ROOT_DIR to include.
26function(declare_mlir_python_sources name)
27  cmake_parse_arguments(ARG
28    ""
29    "ROOT_DIR;ADD_TO_PARENT"
30    "SOURCES;SOURCES_GLOB"
31    ${ARGN})
32
33  if(NOT ARG_ROOT_DIR)
34    set(ARG_ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}")
35  endif()
36
37  # Process the glob.
38  set(_glob_sources)
39  if(ARG_SOURCES_GLOB)
40    set(_glob_spec ${ARG_SOURCES_GLOB})
41    list(TRANSFORM _glob_spec PREPEND "${ARG_ROOT_DIR}/")
42    file(GLOB_RECURSE _glob_sources
43      RELATIVE "${ARG_ROOT_DIR}"
44      ${_glob_spec}
45    )
46    list(APPEND ARG_SOURCES ${_glob_sources})
47  endif()
48
49  # We create a custom target to carry properties and dependencies for
50  # generated sources.
51  add_custom_target(${name})
52  set(_file_depends "${ARG_SOURCES}")
53  list(TRANSFORM _file_depends PREPEND "${ARG_ROOT_DIR}/")
54  set_target_properties(${name} PROPERTIES
55    PYTHON_SOURCES_TYPE pure
56    PYTHON_ROOT_DIR "${ARG_ROOT_DIR}"
57    PYTHON_SOURCES "${ARG_SOURCES}"
58    PYTHON_FILE_DEPENDS "${_file_depends}"
59    PYTHON_DEPENDS ""
60  )
61
62  # Add to parent.
63  if(ARG_ADD_TO_PARENT)
64    set_property(TARGET ${ARG_ADD_TO_PARENT} APPEND PROPERTY PYTHON_DEPENDS ${name})
65  endif()
66endfunction()
67
68# Function: declare_mlir_python_extension
69# Declares a buildable python extension from C++ source files. The built
70# module is considered a python source file and included as everything else.
71# Arguments:
72#   MODULE_NAME: Local import name of the module (i.e. "_mlir").
73#   ADD_TO_PARENT: Same as for declare_mlir_python_sources.
74#   SOURCES: C++ sources making up the module.
75#   PRIVATE_LINK_LIBS: List of libraries to link in privately to the module
76#     regardless of how it is included in the project (generally should be
77#     static libraries that can be included with hidden visibility).
78#   EMBED_CAPI_LINK_LIBS: Dependent CAPI libraries that this extension depends
79#     on. These will be collected for all extensions and put into an
80#     aggregate dylib that is linked against.
81function(declare_mlir_python_extension name)
82  cmake_parse_arguments(ARG
83    ""
84    "MODULE_NAME;ADD_TO_PARENT"
85    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
86    ${ARGN})
87
88  add_custom_target(${name})
89  set_target_properties(${name} PROPERTIES
90    PYTHON_SOURCES_TYPE extension
91    PYTHON_EXTENSION_MODULE_NAME "${ARG_MODULE_NAME}"
92    PYTHON_CPP_SOURCES "${ARG_SOURCES}"
93    PYTHON_PRIVATE_LINK_LIBS "${ARG_PRIVATE_LINK_LIBS}"
94    PYTHON_EMBED_CAPI_LINK_LIBS "${ARG_EMBED_CAPI_LINK_LIBS}"
95    PYTHON_FILE_DEPENDS ""
96    PYTHON_DEPENDS ""
97  )
98
99  # Add to parent.
100  if(ARG_ADD_TO_PARENT)
101    set_property(TARGET ${ARG_ADD_TO_PARENT} APPEND PROPERTY PYTHON_DEPENDS ${name})
102  endif()
103endfunction()
104
105# Function: add_mlir_python_modules
106# Adds python modules to a project, building them from a list of declared
107# source groupings (see declare_mlir_python_sources and
108# declare_mlir_python_extension). One of these must be called for each
109# packaging root in use.
110# Arguments:
111#   ROOT_PREFIX: The directory in the build tree to emit sources. This will
112#     typically be something like ${MY_BINARY_DIR}/python_packages/foobar
113#     for non-relocatable modules or a deeper directory tree for relocatable.
114#   INSTALL_PREFIX: Prefix into the install tree for installing the package.
115#     Typically mirrors the path above but without an absolute path.
116#   DECLARED_SOURCES: List of declared source groups to include. The entire
117#     DAG of source modules is included.
118#   COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every
119#     extension depend on (see mlir_python_add_common_capi_library).
120function(add_mlir_python_modules name)
121  cmake_parse_arguments(ARG
122    ""
123    "ROOT_PREFIX;INSTALL_PREFIX;COMMON_CAPI_LINK_LIBS"
124    "DECLARED_SOURCES"
125    ${ARGN})
126  # Helper to process an individual target.
127  function(_process_target modules_target sources_target)
128    get_target_property(_source_type ${sources_target} PYTHON_SOURCES_TYPE)
129    if(_source_type STREQUAL "pure")
130      # Pure python sources to link into the tree.
131      get_target_property(_python_root_dir ${sources_target} PYTHON_ROOT_DIR)
132      get_target_property(_python_sources ${sources_target} PYTHON_SOURCES)
133      foreach(_source_relative_path ${_python_sources})
134        set(_src_path "${_python_root_dir}/${_source_relative_path}")
135        set(_dest_path "${ARG_ROOT_PREFIX}/${_source_relative_path}")
136
137        get_filename_component(_dest_dir "${_dest_path}" DIRECTORY)
138        get_filename_component(_install_path "${ARG_INSTALL_PREFIX}/${_source_relative_path}" DIRECTORY)
139
140        file(MAKE_DIRECTORY "${_dest_dir}")
141        add_custom_command(
142          TARGET ${modules_target} PRE_BUILD
143          COMMENT "Copying python source ${_src_path} -> ${_dest_path}"
144          DEPENDS "${_src_path}"
145          BYPRODUCTS "${_dest_path}"
146          COMMAND "${CMAKE_COMMAND}" -E create_symlink
147              "${_src_path}" "${_dest_path}"
148        )
149        install(
150          FILES "${_src_path}"
151          DESTINATION "${_install_path}"
152          COMPONENT ${modules_target}
153        )
154      endforeach()
155    elseif(_source_type STREQUAL "extension")
156      # Native CPP extension.
157      get_target_property(_module_name ${sources_target} PYTHON_EXTENSION_MODULE_NAME)
158      get_target_property(_cpp_sources ${sources_target} PYTHON_CPP_SOURCES)
159      get_target_property(_private_link_libs ${sources_target} PYTHON_PRIVATE_LINK_LIBS)
160      set(_extension_target "${name}.extension.${_module_name}.dso")
161      add_mlir_python_extension(${_extension_target} "${_module_name}"
162        INSTALL_COMPONENT ${modules_target}
163        INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
164        OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
165        SOURCES ${_cpp_sources}
166        LINK_LIBS PRIVATE
167          ${_private_link_libs}
168          ${ARG_COMMON_CAPI_LINK_LIBS}
169      )
170      add_dependencies(${name} ${_extension_target})
171      mlir_python_setup_extension_rpath(${_extension_target})
172    else()
173      message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")
174      return()
175    endif()
176  endfunction()
177
178  _flatten_mlir_python_targets(_flat_targets ${ARG_DECLARED_SOURCES})
179  # Collect dependencies.
180  set(_depends)
181  foreach(sources_target ${_flat_targets})
182    get_target_property(_local_depends ${sources_target} PYTHON_FILE_DEPENDS)
183    list(APPEND _depends ${_local_depends})
184  endforeach()
185
186  # Build the modules target.
187  add_custom_target(${name} ALL DEPENDS ${_depends})
188  foreach(sources_target ${_flat_targets})
189    _process_target(${name} ${sources_target})
190  endforeach()
191
192  # Create an install target.
193  if (NOT LLVM_ENABLE_IDE)
194    add_llvm_install_targets(
195      install-${name}
196      DEPENDS ${name}
197      COMPONENT ${name})
198  endif()
199endfunction()
200
201# Function: declare_mlir_dialect_python_bindings
202# Helper to generate source groups for dialects, including both static source
203# files and a TD_FILE to generate wrappers.
204#
205# This will generate a source group named ${ADD_TO_PARENT}.${DIALECT_NAME}.
206#
207# Arguments:
208#   ROOT_DIR: Same as for declare_mlir_python_sources().
209#   ADD_TO_PARENT: Same as for declare_mlir_python_sources(). Unique names
210#     for the subordinate source groups are derived from this.
211#   TD_FILE: Tablegen file to generate source for (relative to ROOT_DIR).
212#   DIALECT_NAME: Python name of the dialect.
213#   SOURCES: Same as declare_mlir_python_sources().
214#   SOURCES_GLOB: Same as declare_mlir_python_sources().
215#   DEPENDS: Additional dependency targets.
216function(declare_mlir_dialect_python_bindings)
217  cmake_parse_arguments(ARG
218    ""
219    "ROOT_DIR;ADD_TO_PARENT;TD_FILE;DIALECT_NAME"
220    "SOURCES;SOURCES_GLOB;DEPENDS"
221    ${ARGN})
222  # Sources.
223  set(_dialect_target "${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}")
224  declare_mlir_python_sources(${_dialect_target}
225    ROOT_DIR "${ARG_ROOT_DIR}"
226    ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
227    SOURCES "${ARG_SOURCES}"
228    SOURCES_GLOB "${ARG_SOURCES_GLOB}"
229  )
230
231  # Tablegen
232  if(ARG_TD_FILE)
233    set(tblgen_target "${ARG_ADD_TO}.${ARG_DIALECT_NAME}.tablegen")
234    set(td_file "${ARG_ROOT_DIR}/${ARG_TD_FILE}")
235    get_filename_component(relative_td_directory "${ARG_TD_FILE}" DIRECTORY)
236    set(dialect_filename "${relative_td_directory}/_${ARG_DIALECT_NAME}_ops_gen.py")
237    set(LLVM_TARGET_DEFINITIONS ${td_file})
238    mlir_tablegen("${dialect_filename}" -gen-python-op-bindings
239                  -bind-dialect=${ARG_DIALECT_NAME})
240    add_public_tablegen_target(${tblgen_target})
241    if(ARG_DEPENDS)
242      add_dependencies(${tblgen_target} ${ARG_DEPENDS})
243    endif()
244
245    # Generated.
246    declare_mlir_python_sources("${ARG_ADD_TO_PARENT}.${ARG_DIALECT_NAME}.ops_gen"
247      ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
248      ADD_TO_PARENT "${_dialect_target}"
249      SOURCES "${dialect_filename}"
250    )
251  endif()
252endfunction()
253
254# Function: mlir_python_setup_extension_rpath
255# Sets RPATH properties on a target, assuming that it is being output to
256# an _mlir_libs directory with all other libraries. For static linkage,
257# the RPATH will just be the origin. If linking dynamically, then the LLVM
258# library directory will be added.
259# Arguments:
260#   RELATIVE_INSTALL_ROOT: If building dynamically, an RPATH entry will be
261#     added to the install tree lib/ directory by first traversing this
262#     path relative to the installation location. Typically a number of ".."
263#     entries, one for each level of the install path.
264function(mlir_python_setup_extension_rpath target)
265  cmake_parse_arguments(ARG
266    ""
267    "RELATIVE_INSTALL_ROOT"
268    ""
269    ${ARGN})
270
271  # RPATH handling.
272  # For the build tree, include the LLVM lib directory and the current
273  # directory for RPATH searching. For install, just the current directory
274  # (assumes that needed dependencies have been installed).
275  if(NOT APPLE AND NOT UNIX)
276    return()
277  endif()
278
279  set(_origin_prefix "\$ORIGIN")
280  if(APPLE)
281    set(_origin_prefix "@loader_path")
282  endif()
283  set_target_properties(${target} PROPERTIES
284    BUILD_WITH_INSTALL_RPATH OFF
285    BUILD_RPATH "${_origin_prefix}"
286    INSTALL_RPATH "${_origin_prefix}"
287  )
288
289  # For static builds, that is all that is needed: all dependencies will be in
290  # the one directory. For shared builds, then we also need to add the global
291  # lib directory. This will be absolute for the build tree and relative for
292  # install.
293  # When we have access to CMake >= 3.20, there is a helper to calculate this.
294  if(BUILD_SHARED_LIBS AND ARG_RELATIVE_INSTALL_ROOT)
295    get_filename_component(_real_lib_dir "${LLVM_LIBRARY_OUTPUT_INTDIR}" REALPATH)
296    set_property(TARGET ${target} APPEND PROPERTY
297      BUILD_RPATH "${_real_lib_dir}")
298    set_property(TARGET ${target} APPEND PROPERTY
299      INSTALL_RPATH "${_origin_prefix}/${ARG_RELATIVE_INSTALL_ROOT}/lib${LLVM_LIBDIR_SUFFIX}")
300  endif()
301endfunction()
302
303# Function: add_mlir_python_common_capi_library
304# Adds a shared library which embeds dependent CAPI libraries needed to link
305# all extensions.
306# Arguments:
307#   INSTALL_COMPONENT: Name of the install component. Typically same as the
308#     target name passed to add_mlir_python_modules().
309#   INSTALL_DESTINATION: Prefix into the install tree in which to install the
310#     library.
311#   OUTPUT_DIRECTORY: Full path in the build tree in which to create the
312#     library. Typically, this will be the common _mlir_libs directory where
313#     all extensions are emitted.
314#   RELATIVE_INSTALL_ROOT: See mlir_python_setup_extension_rpath().
315#   DECLARED_SOURCES: Source groups from which to discover dependent
316#     EMBED_CAPI_LINK_LIBS.
317#   EMBED_LIBS: Additional libraries to embed (must be built with OBJECTS and
318#     have an "obj.${name}" object library associated).
319function(add_mlir_python_common_capi_library name)
320  cmake_parse_arguments(ARG
321    ""
322    "INSTALL_COMPONENT;INSTALL_DESTINATION;OUTPUT_DIRECTORY;RELATIVE_INSTALL_ROOT"
323    "DECLARED_SOURCES;EMBED_LIBS"
324    ${ARGN})
325  # TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
326  # once ready.
327
328  # Collect all explicit and transitive embed libs.
329  set(_embed_libs ${ARG_EMBED_LIBS})
330  _flatten_mlir_python_targets(_all_source_targets ${ARG_DECLARED_SOURCES})
331  foreach(t ${_all_source_targets})
332    get_target_property(_local_embed_libs ${t} PYTHON_EMBED_CAPI_LINK_LIBS)
333    if(_local_embed_libs)
334      list(APPEND _embed_libs ${_local_embed_libs})
335    endif()
336  endforeach()
337  list(REMOVE_DUPLICATES _embed_libs)
338
339  foreach(lib ${_embed_libs})
340    if(XCODE)
341      # Xcode doesn't support object libraries, so we have to trick it into
342      # linking the static libraries instead.
343      list(APPEND _deps "-force_load" ${lib})
344    else()
345      list(APPEND _objects $<TARGET_OBJECTS:obj.${lib}>)
346    endif()
347    # Accumulate transitive deps of each exported lib into _DEPS.
348    list(APPEND _deps $<TARGET_PROPERTY:${lib},LINK_LIBRARIES>)
349  endforeach()
350
351  add_mlir_library(${name}
352    PARTIAL_SOURCES_INTENDED
353    SHARED
354    DISABLE_INSTALL
355    ${_objects}
356    EXCLUDE_FROM_LIBMLIR
357    LINK_LIBS
358    ${_deps}
359  )
360  if(MSVC)
361    set_property(TARGET ${name} PROPERTY WINDOWS_EXPORT_ALL_SYMBOLS ON)
362  endif()
363  set_target_properties(${name} PROPERTIES
364    LIBRARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
365    BINARY_OUTPUT_DIRECTORY "${ARG_OUTPUT_DIRECTORY}"
366  )
367  mlir_python_setup_extension_rpath(${name}
368    RELATIVE_INSTALL_ROOT "${ARG_RELATIVE_INSTALL_ROOT}"
369  )
370  install(TARGETS ${name}
371    COMPONENT ${ARG_INSTALL_COMPONENT}
372    LIBRARY DESTINATION "${ARG_INSTALL_DESTINATION}"
373    RUNTIME DESTINATION "${ARG_INSTALL_DESTINATION}"
374  )
375
376endfunction()
377
378function(_flatten_mlir_python_targets output_var)
379  set(_flattened)
380  foreach(t ${ARGN})
381    get_target_property(_source_type ${t} PYTHON_SOURCES_TYPE)
382    get_target_property(_depends ${t} PYTHON_DEPENDS)
383    if(_source_type)
384      list(APPEND _flattened "${t}")
385      if(_depends)
386        _flatten_mlir_python_targets(_local_flattened ${_depends})
387        list(APPEND _flattened ${_local_flattened})
388      endif()
389    endif()
390  endforeach()
391  list(REMOVE_DUPLICATES _flattened)
392  set(${output_var} "${_flattened}" PARENT_SCOPE)
393endfunction()
394
395################################################################################
396# Build python extension
397################################################################################
398function(add_mlir_python_extension libname extname)
399  cmake_parse_arguments(ARG
400  ""
401  "INSTALL_COMPONENT;INSTALL_DIR;OUTPUT_DIRECTORY"
402  "SOURCES;LINK_LIBS"
403  ${ARGN})
404  if (ARG_UNPARSED_ARGUMENTS)
405    message(FATAL_ERROR " Unhandled arguments to add_mlir_python_extension(${libname}, ... : ${ARG_UNPARSED_ARGUMENTS}")
406  endif()
407  if ("${ARG_SOURCES}" STREQUAL "")
408    message(FATAL_ERROR " Missing SOURCES argument to add_mlir_python_extension(${libname}, ...")
409  endif()
410
411  # Build-time RPath layouts require to be a directory one up from the
412  # binary root.
413  # TODO: Don't reference the LLVM_BINARY_DIR here: the invariant is that
414  # the output directory must be at the same level of the lib directory
415  # where libMLIR.so is installed. This is presently not optimal from a
416  # project separation perspective and a discussion on how to better
417  # segment MLIR libraries needs to happen.
418  # TODO: Remove this when downstreams are moved off of it.
419  if(NOT ARG_OUTPUT_DIRECTORY)
420    set(ARG_OUTPUT_DIRECTORY ${LLVM_BINARY_DIR}/python)
421  endif()
422
423  # Normally on unix-like platforms, extensions are built as "MODULE" libraries
424  # and do not explicitly link to the python shared object. This allows for
425  # some greater deployment flexibility since the extension will bind to
426  # symbols in the python interpreter on load. However, it also keeps the
427  # linker from erroring on undefined symbols, leaving this to (usually obtuse)
428  # runtime errors. Building in "SHARED" mode with an explicit link to the
429  # python libraries allows us to build with the expectation of no undefined
430  # symbols, which is better for development. Note that not all python
431  # configurations provide build-time libraries to link against, in which
432  # case, we fall back to MODULE linking.
433  if(Python3_LIBRARIES STREQUAL "" OR NOT MLIR_BINDINGS_PYTHON_LOCK_VERSION)
434    set(PYEXT_LINK_MODE MODULE)
435    set(PYEXT_LIBADD)
436  else()
437    set(PYEXT_LINK_MODE SHARED)
438    set(PYEXT_LIBADD ${Python3_LIBRARIES})
439  endif()
440
441  # The actual extension library produces a shared-object or DLL and has
442  # sources that must be compiled in accordance with pybind11 needs (RTTI and
443  # exceptions).
444  add_library(${libname}
445    ${PYEXT_LINK_MODE}
446    ${ARG_SOURCES}
447  )
448
449  target_include_directories(${libname} PRIVATE
450    "${Python3_INCLUDE_DIRS}"
451    "${pybind11_INCLUDE_DIR}"
452  )
453
454  target_link_directories(${libname} PRIVATE
455    "${Python3_LIBRARY_DIRS}"
456  )
457
458  # The extension itself must be compiled with RTTI and exceptions enabled.
459  # Also, some warning classes triggered by pybind11 are disabled.
460  target_compile_options(${libname} PRIVATE
461    $<$<OR:$<CXX_COMPILER_ID:Clang>,$<CXX_COMPILER_ID:AppleClang>,$<CXX_COMPILER_ID:GNU>>:
462      # Enable RTTI and exceptions.
463      -frtti -fexceptions
464      # Noisy pybind warnings
465      -Wno-unused-value
466      -Wno-covered-switch-default
467    >
468    $<$<CXX_COMPILER_ID:MSVC>:
469      # Enable RTTI and exceptions.
470      /EHsc /GR>
471  )
472
473  # Configure the output to match python expectations.
474  set_target_properties(
475    ${libname} PROPERTIES
476    LIBRARY_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
477    OUTPUT_NAME "${extname}"
478    PREFIX "${PYTHON_MODULE_PREFIX}"
479    SUFFIX "${PYTHON_MODULE_SUFFIX}${PYTHON_MODULE_EXTENSION}"
480  )
481
482  if(WIN32)
483    # Need to also set the RUNTIME_OUTPUT_DIRECTORY on Windows in order to
484    # control where the .dll gets written.
485    set_target_properties(
486      ${libname} PROPERTIES
487      RUNTIME_OUTPUT_DIRECTORY ${ARG_OUTPUT_DIRECTORY}
488    )
489  endif()
490
491  # pybind11 requires binding code to be compiled with -fvisibility=hidden
492  # For static linkage, better code can be generated if the entire project
493  # compiles that way, but that is not enforced here. Instead, include a linker
494  # script that explicitly hides anything but the PyInit_* symbols, allowing gc
495  # to take place.
496  set_target_properties(${libname} PROPERTIES CXX_VISIBILITY_PRESET "hidden")
497
498  # Python extensions depends *only* on the public API and LLVMSupport unless
499  # if further dependencies are added explicitly.
500  target_link_libraries(${libname}
501    PRIVATE
502    ${ARG_LINK_LIBS}
503    ${PYEXT_LIBADD}
504  )
505
506  target_link_options(${libname}
507    PRIVATE
508      # On Linux, disable re-export of any static linked libraries that
509      # came through.
510      $<$<PLATFORM_ID:Linux>:LINKER:--exclude-libs,ALL>
511  )
512
513  ################################################################################
514  # Install
515  ################################################################################
516  if (ARG_INSTALL_DIR)
517    install(TARGETS ${libname}
518      COMPONENT ${ARG_INSTALL_COMPONENT}
519      LIBRARY DESTINATION ${ARG_INSTALL_DIR}
520      ARCHIVE DESTINATION ${ARG_INSTALL_DIR}
521      # NOTE: Even on DLL-platforms, extensions go in the lib directory tree.
522      RUNTIME DESTINATION ${ARG_INSTALL_DIR}
523    )
524  endif()
525endfunction()
526