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