1include(AddMLIRPython)
2
3################################################################################
4# Structural groupings.
5################################################################################
6
7declare_mlir_python_sources(MLIRPythonSources)
8declare_mlir_python_sources(MLIRPythonSources.Dialects
9  ADD_TO_PARENT MLIRPythonSources)
10
11################################################################################
12# Pure python sources and generated code
13################################################################################
14
15declare_mlir_python_sources(MLIRPythonSources.Core
16  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
17  ADD_TO_PARENT MLIRPythonSources
18  SOURCES
19    _mlir_libs/__init__.py
20    ir.py
21    passmanager.py
22    dialects/_ods_common.py
23
24    # The main _mlir module has submodules: include stubs from each.
25    _mlir_libs/_mlir/__init__.pyi
26    _mlir_libs/_mlir/ir.pyi
27    _mlir_libs/_mlir/passmanager.pyi
28)
29
30declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
31  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
32  ADD_TO_PARENT MLIRPythonSources
33  SOURCES
34    execution_engine.py
35    _mlir_libs/_mlirExecutionEngine.pyi
36  SOURCES_GLOB
37    runtime/*.py
38)
39
40declare_mlir_python_sources(MLIRPythonSources.Passes
41  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
42  ADD_TO_PARENT MLIRPythonSources
43  SOURCES_GLOB
44    all_passes_registration/*.py
45    conversions/*.py
46    transforms/*.py
47)
48
49declare_mlir_python_sources(MLIRPythonCAPIHeaderSources
50  ROOT_DIR "${MLIR_SOURCE_DIR}/include"
51  SOURCES_GLOB "mlir-c/*.h"
52  DEST_PREFIX "_mlir_libs/include"
53)
54
55################################################################################
56# Dialect bindings
57################################################################################
58
59declare_mlir_dialect_python_bindings(
60  ADD_TO_PARENT MLIRPythonSources.Dialects
61  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
62  TD_FILE dialects/AsyncOps.td
63  SOURCES_GLOB dialects/async_dialect/*.py
64  DIALECT_NAME async_dialect)
65
66declare_mlir_dialect_python_bindings(
67  ADD_TO_PARENT MLIRPythonSources.Dialects
68  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
69  TD_FILE dialects/BuiltinOps.td
70  SOURCES
71    dialects/builtin.py
72    dialects/_builtin_ops_ext.py
73  DIALECT_NAME builtin)
74
75declare_mlir_dialect_python_bindings(
76  ADD_TO_PARENT MLIRPythonSources.Dialects
77  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
78  TD_FILE dialects/GPUOps.td
79  SOURCES_GLOB dialects/gpu/*.py
80  DIALECT_NAME gpu)
81
82declare_mlir_dialect_python_bindings(
83  ADD_TO_PARENT MLIRPythonSources.Dialects
84  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
85  TD_FILE dialects/LinalgOps.td
86  SOURCES
87    dialects/_linalg_ops_ext.py
88  SOURCES_GLOB
89    dialects/linalg/*.py
90  DIALECT_NAME linalg
91  DEPENDS LinalgOdsGen)
92
93declare_mlir_dialect_python_bindings(
94  ADD_TO_PARENT MLIRPythonSources.Dialects
95  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
96  TD_FILE dialects/MathOps.td
97  SOURCES dialects/math.py
98  DIALECT_NAME math)
99
100declare_mlir_dialect_python_bindings(
101  ADD_TO_PARENT MLIRPythonSources.Dialects
102  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
103  TD_FILE dialects/ArithmeticOps.td
104  SOURCES
105    dialects/arith.py
106    dialects/_arith_ops_ext.py
107  DIALECT_NAME arith)
108
109declare_mlir_dialect_python_bindings(
110  ADD_TO_PARENT MLIRPythonSources.Dialects
111  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
112  TD_FILE dialects/MemRefOps.td
113  SOURCES
114    dialects/memref.py
115    dialects/_memref_ops_ext.py
116  DIALECT_NAME memref)
117
118declare_mlir_dialect_python_bindings(
119  ADD_TO_PARENT MLIRPythonSources.Dialects
120  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
121  TD_FILE dialects/SCFOps.td
122  SOURCES
123    dialects/scf.py
124    dialects/_scf_ops_ext.py
125  DIALECT_NAME scf)
126
127declare_mlir_dialect_python_bindings(
128  ADD_TO_PARENT MLIRPythonSources.Dialects
129  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
130  TD_FILE dialects/ShapeOps.td
131  SOURCES dialects/shape.py
132  DIALECT_NAME shape)
133
134declare_mlir_dialect_python_bindings(
135  ADD_TO_PARENT MLIRPythonSources.Dialects
136  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
137  TD_FILE dialects/SparseTensorOps.td
138  SOURCES dialects/sparse_tensor.py
139  DIALECT_NAME sparse_tensor)
140
141declare_mlir_dialect_python_bindings(
142  ADD_TO_PARENT MLIRPythonSources.Dialects
143  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
144  TD_FILE dialects/StandardOps.td
145  SOURCES
146    dialects/std.py
147    dialects/_std_ops_ext.py
148  DIALECT_NAME std)
149
150declare_mlir_dialect_python_bindings(
151  ADD_TO_PARENT MLIRPythonSources.Dialects
152  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
153  TD_FILE dialects/TensorOps.td
154  SOURCES dialects/tensor.py
155  DIALECT_NAME tensor)
156
157declare_mlir_dialect_python_bindings(
158  ADD_TO_PARENT MLIRPythonSources.Dialects
159  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
160  TD_FILE dialects/TosaOps.td
161  SOURCES dialects/tosa.py
162  DIALECT_NAME tosa)
163
164declare_mlir_dialect_python_bindings(
165  ADD_TO_PARENT MLIRPythonSources.Dialects
166  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
167  TD_FILE dialects/VectorOps.td
168  SOURCES dialects/vector.py
169  DIALECT_NAME vector)
170
171################################################################################
172# Python extensions.
173# The sources for these are all in lib/Bindings/Python, but since they have to
174# be rebuilt for each package and integrate with the source setup here, we
175# just reference them here instead of having ordered, cross package target
176# dependencies.
177################################################################################
178
179set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
180declare_mlir_python_extension(MLIRPythonExtension.Core
181  MODULE_NAME _mlir
182  ADD_TO_PARENT MLIRPythonSources.Core
183  ROOT_DIR "${PYTHON_SOURCE_DIR}"
184  SOURCES
185    DialectLinalg.cpp  # TODO: Break this out.
186    DialectSparseTensor.cpp  # TODO: Break this out.
187    MainModule.cpp
188    IRAffine.cpp
189    IRAttributes.cpp
190    IRCore.cpp
191    IRInterfaces.cpp
192    IRModule.cpp
193    IRTypes.cpp
194    PybindUtils.cpp
195    Pass.cpp
196
197    # Headers must be included explicitly so they are installed.
198    Dialects.h
199    Globals.h
200    IRModule.h
201    Pass.h
202    PybindUtils.h
203  PRIVATE_LINK_LIBS
204    LLVMSupport
205  EMBED_CAPI_LINK_LIBS
206    MLIRCAPIDebug
207    MLIRCAPIIR
208    MLIRCAPIInterfaces
209    MLIRCAPIRegistration  # TODO: See about dis-aggregating
210
211    # Dialects
212    MLIRCAPILinalg  # TODO: Remove when above is removed.
213    MLIRCAPISparseTensor  # TODO: Remove when above is removed.
214    MLIRCAPIStandard
215)
216
217declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
218  MODULE_NAME _mlirAllPassesRegistration
219  ROOT_DIR "${PYTHON_SOURCE_DIR}"
220  SOURCES
221    AllPassesRegistration.cpp
222  PRIVATE_LINK_LIBS
223    LLVMSupport
224  EMBED_CAPI_LINK_LIBS
225    MLIRCAPIConversion
226    MLIRCAPITransforms
227)
228
229declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
230  MODULE_NAME _mlirAsyncPasses
231  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
232  ROOT_DIR "${PYTHON_SOURCE_DIR}"
233  SOURCES
234    AsyncPasses.cpp
235  PRIVATE_LINK_LIBS
236    LLVMSupport
237  EMBED_CAPI_LINK_LIBS
238    MLIRCAPIAsync
239)
240
241declare_mlir_python_extension(MLIRPythonExtension.Conversions
242  MODULE_NAME _mlirConversions
243  ADD_TO_PARENT MLIRPythonSources.Passes
244  ROOT_DIR "${PYTHON_SOURCE_DIR}"
245  SOURCES
246    Conversions/Conversions.cpp
247  PRIVATE_LINK_LIBS
248    LLVMSupport
249  EMBED_CAPI_LINK_LIBS
250  MLIRCAPIConversion
251)
252
253declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
254  MODULE_NAME _mlirExecutionEngine
255  ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
256  ROOT_DIR "${PYTHON_SOURCE_DIR}"
257  SOURCES
258    ExecutionEngineModule.cpp
259  PRIVATE_LINK_LIBS
260    LLVMSupport
261  EMBED_CAPI_LINK_LIBS
262    MLIRCAPIExecutionEngine
263)
264
265declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
266  MODULE_NAME _mlirGPUPasses
267  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
268  ROOT_DIR "${PYTHON_SOURCE_DIR}"
269  SOURCES
270    GPUPasses.cpp
271  PRIVATE_LINK_LIBS
272    LLVMSupport
273  EMBED_CAPI_LINK_LIBS
274    MLIRCAPIGPU
275)
276
277declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
278  MODULE_NAME _mlirLinalgPasses
279  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
280  ROOT_DIR "${PYTHON_SOURCE_DIR}"
281  SOURCES
282    LinalgPasses.cpp
283  PRIVATE_LINK_LIBS
284    LLVMSupport
285  EMBED_CAPI_LINK_LIBS
286    MLIRCAPILinalg
287)
288
289declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
290  MODULE_NAME _mlirSparseTensorPasses
291  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
292  ROOT_DIR "${PYTHON_SOURCE_DIR}"
293  SOURCES
294    SparseTensorPasses.cpp
295  PRIVATE_LINK_LIBS
296    LLVMSupport
297  EMBED_CAPI_LINK_LIBS
298    MLIRCAPISparseTensor
299)
300
301declare_mlir_python_extension(MLIRPythonExtension.Transforms
302  MODULE_NAME _mlirTransforms
303  ADD_TO_PARENT MLIRPythonSources.Passes
304  ROOT_DIR "${PYTHON_SOURCE_DIR}"
305  SOURCES
306    Transforms/Transforms.cpp
307  PRIVATE_LINK_LIBS
308    LLVMSupport
309  EMBED_CAPI_LINK_LIBS
310    MLIRCAPITransforms
311)
312
313# TODO: Figure out how to put this in the test tree.
314# This should not be included in the main Python extension. However,
315# putting it into MLIRPythonTestSources along with the dialect declaration
316# above confuses Python module loader when running under lit.
317set(_ADDL_TEST_SOURCES)
318if(MLIR_INCLUDE_TESTS)
319  set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
320  declare_mlir_python_sources(MLIRPythonTestSources)
321  declare_mlir_python_sources(MLIRPythonTestSources.Dialects
322    ADD_TO_PARENT MLIRPythonTestSources)
323
324  # TODO: this uses a tablegen file from the test directory and should be
325  # decoupled from here.
326  declare_mlir_python_sources(
327    MLIRPythonTestSources.Dialects.PythonTest
328    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
329    ADD_TO_PARENT MLIRPythonTestSources.Dialects
330    SOURCES dialects/python_test.py)
331  set(LLVM_TARGET_DEFINITIONS
332    "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
333  mlir_tablegen(
334    "dialects/_python_test_ops_gen.py"
335    -gen-python-op-bindings
336    -bind-dialect=python_test)
337  add_public_tablegen_target(PythonTestDialectPyIncGen)
338  declare_mlir_python_sources(
339    MLIRPythonTestSources.Dialects.PythonTest.ops_gen
340    ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
341    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
342    SOURCES "dialects/_python_test_ops_gen.py")
343
344  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
345    MODULE_NAME _mlirPythonTest
346    ADD_TO_PARENT MLIRPythonTestSources.Dialects
347    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
348    SOURCES
349      PythonTestModule.cpp
350    PRIVATE_LINK_LIBS
351      LLVMSupport
352    EMBED_CAPI_LINK_LIBS
353      MLIRCAPIPythonTestDialect
354  )
355endif()
356
357################################################################################
358# Common CAPI dependency DSO.
359# All python extensions must link through one DSO which exports the CAPI, and
360# this must have a globally unique name amongst all embeddors of the python
361# library since it will effectively have global scope.
362#
363# The presence of this aggregate library is part of the long term plan, but its
364# use needs to be made more flexible.
365#
366# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
367# once ready.
368################################################################################
369
370add_mlir_python_common_capi_library(MLIRPythonCAPI
371  INSTALL_COMPONENT MLIRPythonModules
372  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
373  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
374  RELATIVE_INSTALL_ROOT "../../../.."
375  DECLARED_SOURCES
376    MLIRPythonSources
377    MLIRPythonExtension.AllPassesRegistration
378    ${_ADDL_TEST_SOURCES}
379)
380
381################################################################################
382# The fully assembled package of modules.
383# This must come last.
384################################################################################
385
386add_mlir_python_modules(MLIRPythonModules
387  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
388  INSTALL_PREFIX "python_packages/mlir_core/mlir"
389  DECLARED_SOURCES
390    MLIRPythonSources
391    MLIRPythonExtension.AllPassesRegistration
392    MLIRPythonCAPIHeaderSources
393    ${_ADDL_TEST_SOURCES}
394  COMMON_CAPI_LINK_LIBS
395    MLIRPythonCAPI
396  )
397