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(MLIRPythonCAPI.HeaderSources
50  ROOT_DIR "${MLIR_MAIN_INCLUDE_DIR}"
51  SOURCES_GLOB "mlir-c/*.h"
52)
53
54################################################################################
55# Dialect bindings
56################################################################################
57
58declare_mlir_dialect_python_bindings(
59  ADD_TO_PARENT MLIRPythonSources.Dialects
60  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
61  TD_FILE dialects/AsyncOps.td
62  SOURCES_GLOB dialects/async_dialect/*.py
63  DIALECT_NAME async_dialect)
64
65declare_mlir_dialect_python_bindings(
66  ADD_TO_PARENT MLIRPythonSources.Dialects
67  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
68  TD_FILE dialects/BufferizationOps.td
69  SOURCES
70    dialects/bufferization.py
71    dialects/_bufferization_ops_ext.py
72  DIALECT_NAME bufferization)
73
74declare_mlir_dialect_python_bindings(
75  ADD_TO_PARENT MLIRPythonSources.Dialects
76  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
77  TD_FILE dialects/BuiltinOps.td
78  SOURCES
79    dialects/builtin.py
80    dialects/_builtin_ops_ext.py
81  DIALECT_NAME builtin)
82
83declare_mlir_dialect_python_bindings(
84  ADD_TO_PARENT MLIRPythonSources.Dialects
85  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
86  TD_FILE dialects/ComplexOps.td
87  SOURCES
88    dialects/complex.py
89  DIALECT_NAME complex)
90
91declare_mlir_dialect_python_bindings(
92  ADD_TO_PARENT MLIRPythonSources.Dialects
93  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
94  TD_FILE dialects/ControlFlowOps.td
95  SOURCES
96    dialects/cf.py
97  DIALECT_NAME cf)
98
99declare_mlir_dialect_python_bindings(
100  ADD_TO_PARENT MLIRPythonSources.Dialects
101  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
102  TD_FILE dialects/FuncOps.td
103  SOURCES
104    dialects/func.py
105    dialects/_func_ops_ext.py
106  DIALECT_NAME func)
107
108declare_mlir_dialect_python_bindings(
109  ADD_TO_PARENT MLIRPythonSources.Dialects
110  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
111  TD_FILE dialects/GPUOps.td
112  SOURCES_GLOB dialects/gpu/*.py
113  DIALECT_NAME gpu)
114
115declare_mlir_dialect_python_bindings(
116  ADD_TO_PARENT MLIRPythonSources.Dialects
117  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
118  TD_FILE dialects/LinalgOps.td
119  SOURCES
120    dialects/_linalg_ops_ext.py
121  SOURCES_GLOB
122    dialects/linalg/*.py
123  DIALECT_NAME linalg
124  DEPENDS LinalgOdsGen)
125
126declare_mlir_dialect_python_bindings(
127  ADD_TO_PARENT MLIRPythonSources.Dialects
128  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
129  TD_FILE dialects/TransformOps.td
130  SOURCES
131    dialects/_transform_ops_ext.py
132    dialects/transform/__init__.py
133  DIALECT_NAME transform)
134
135declare_mlir_dialect_extension_python_bindings(
136  ADD_TO_PARENT MLIRPythonSources.Dialects
137  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
138  TD_FILE dialects/SCFLoopTransformOps.td
139  SOURCES
140    dialects/_loop_transform_ops_ext.py
141    dialects/transform/loop.py
142  DIALECT_NAME transform
143  EXTENSION_NAME loop_transform)
144
145declare_mlir_dialect_extension_python_bindings(
146  ADD_TO_PARENT MLIRPythonSources.Dialects
147  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
148  TD_FILE dialects/LinalgStructuredTransformOps.td
149  SOURCES
150    dialects/_structured_transform_ops_ext.py
151    dialects/transform/structured.py
152  DIALECT_NAME transform
153  EXTENSION_NAME structured_transform)
154
155declare_mlir_dialect_python_bindings(
156  ADD_TO_PARENT MLIRPythonSources.Dialects
157  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
158  TD_FILE dialects/MathOps.td
159  SOURCES dialects/math.py
160  DIALECT_NAME math)
161
162declare_mlir_dialect_python_bindings(
163  ADD_TO_PARENT MLIRPythonSources.Dialects
164  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
165  TD_FILE dialects/ArithmeticOps.td
166  SOURCES
167    dialects/arith.py
168    dialects/_arith_ops_ext.py
169  DIALECT_NAME arith)
170
171declare_mlir_dialect_python_bindings(
172  ADD_TO_PARENT MLIRPythonSources.Dialects
173  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
174  TD_FILE dialects/MemRefOps.td
175  SOURCES
176    dialects/memref.py
177    dialects/_memref_ops_ext.py
178  DIALECT_NAME memref)
179
180declare_mlir_dialect_python_bindings(
181  ADD_TO_PARENT MLIRPythonSources.Dialects
182  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
183  TD_FILE dialects/MLProgramOps.td
184  SOURCES
185    dialects/ml_program.py
186    dialects/_ml_program_ops_ext.py
187  DIALECT_NAME ml_program)
188
189declare_mlir_python_sources(
190  MLIRPythonSources.Dialects.quant
191  ADD_TO_PARENT MLIRPythonSources.Dialects
192  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
193  SOURCES
194    dialects/quant.py
195    _mlir_libs/_mlir/dialects/quant.pyi)
196
197declare_mlir_dialect_python_bindings(
198  ADD_TO_PARENT MLIRPythonSources.Dialects
199  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
200  TD_FILE dialects/PDLOps.td
201  SOURCES
202    dialects/pdl.py
203    dialects/_pdl_ops_ext.py
204    _mlir_libs/_mlir/dialects/pdl.pyi
205  DIALECT_NAME pdl)
206
207declare_mlir_dialect_python_bindings(
208  ADD_TO_PARENT MLIRPythonSources.Dialects
209  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
210  TD_FILE dialects/SCFOps.td
211  SOURCES
212    dialects/scf.py
213    dialects/_scf_ops_ext.py
214  DIALECT_NAME scf)
215
216declare_mlir_dialect_python_bindings(
217  ADD_TO_PARENT MLIRPythonSources.Dialects
218  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
219  TD_FILE dialects/ShapeOps.td
220  SOURCES dialects/shape.py
221  DIALECT_NAME shape)
222
223declare_mlir_dialect_python_bindings(
224  ADD_TO_PARENT MLIRPythonSources.Dialects
225  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
226  TD_FILE dialects/SparseTensorOps.td
227  SOURCES dialects/sparse_tensor.py
228  DIALECT_NAME sparse_tensor)
229
230declare_mlir_dialect_python_bindings(
231  ADD_TO_PARENT MLIRPythonSources.Dialects
232  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
233  TD_FILE dialects/TensorOps.td
234  SOURCES dialects/tensor.py
235  DIALECT_NAME tensor)
236
237declare_mlir_dialect_python_bindings(
238  ADD_TO_PARENT MLIRPythonSources.Dialects
239  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
240  TD_FILE dialects/TosaOps.td
241  SOURCES dialects/tosa.py
242  DIALECT_NAME tosa)
243
244declare_mlir_dialect_python_bindings(
245  ADD_TO_PARENT MLIRPythonSources.Dialects
246  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
247  TD_FILE dialects/VectorOps.td
248  SOURCES dialects/vector.py
249  DIALECT_NAME vector)
250
251################################################################################
252# Python extensions.
253# The sources for these are all in lib/Bindings/Python, but since they have to
254# be rebuilt for each package and integrate with the source setup here, we
255# just reference them here instead of having ordered, cross package target
256# dependencies.
257################################################################################
258
259set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
260declare_mlir_python_extension(MLIRPythonExtension.Core
261  MODULE_NAME _mlir
262  ADD_TO_PARENT MLIRPythonSources.Core
263  ROOT_DIR "${PYTHON_SOURCE_DIR}"
264  SOURCES
265    MainModule.cpp
266    IRAffine.cpp
267    IRAttributes.cpp
268    IRCore.cpp
269    IRInterfaces.cpp
270    IRModule.cpp
271    IRTypes.cpp
272    PybindUtils.cpp
273    Pass.cpp
274
275    # Headers must be included explicitly so they are installed.
276    Globals.h
277    IRModule.h
278    Pass.h
279    PybindUtils.h
280  PRIVATE_LINK_LIBS
281    LLVMSupport
282  EMBED_CAPI_LINK_LIBS
283    MLIRCAPIDebug
284    MLIRCAPIIR
285    MLIRCAPIInterfaces
286    MLIRCAPIRegistration  # TODO: See about dis-aggregating
287
288    # Dialects
289    MLIRCAPIFunc
290)
291
292declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
293  MODULE_NAME _mlirDialectsLinalg
294  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
295  ROOT_DIR "${PYTHON_SOURCE_DIR}"
296  SOURCES
297    DialectLinalg.cpp
298  PRIVATE_LINK_LIBS
299    LLVMSupport
300  EMBED_CAPI_LINK_LIBS
301    MLIRCAPIIR
302    MLIRCAPILinalg
303)
304
305declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
306  MODULE_NAME _mlirDialectsQuant
307  ADD_TO_PARENT MLIRPythonSources.Dialects.quant
308  ROOT_DIR "${PYTHON_SOURCE_DIR}"
309  SOURCES
310    DialectQuant.cpp
311  PRIVATE_LINK_LIBS
312    LLVMSupport
313  EMBED_CAPI_LINK_LIBS
314    MLIRCAPIIR
315    MLIRCAPIQuant
316)
317
318declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
319  MODULE_NAME _mlirDialectsPDL
320  ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
321  ROOT_DIR "${PYTHON_SOURCE_DIR}"
322  SOURCES
323    DialectPDL.cpp
324  PRIVATE_LINK_LIBS
325    LLVMSupport
326  EMBED_CAPI_LINK_LIBS
327    MLIRCAPIIR
328    MLIRCAPIPDL
329)
330
331declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
332  MODULE_NAME _mlirDialectsSparseTensor
333  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
334  ROOT_DIR "${PYTHON_SOURCE_DIR}"
335  SOURCES
336    DialectSparseTensor.cpp
337  PRIVATE_LINK_LIBS
338    LLVMSupport
339  EMBED_CAPI_LINK_LIBS
340    MLIRCAPIIR
341    MLIRCAPISparseTensor
342)
343
344declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
345  MODULE_NAME _mlirAllPassesRegistration
346  ROOT_DIR "${PYTHON_SOURCE_DIR}"
347  SOURCES
348    AllPassesRegistration.cpp
349  PRIVATE_LINK_LIBS
350    LLVMSupport
351  EMBED_CAPI_LINK_LIBS
352    MLIRCAPIConversion
353    MLIRCAPITransforms
354)
355
356declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
357  MODULE_NAME _mlirAsyncPasses
358  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
359  ROOT_DIR "${PYTHON_SOURCE_DIR}"
360  SOURCES
361    AsyncPasses.cpp
362  PRIVATE_LINK_LIBS
363    LLVMSupport
364  EMBED_CAPI_LINK_LIBS
365    MLIRCAPIAsync
366)
367
368declare_mlir_python_extension(MLIRPythonExtension.Conversions
369  MODULE_NAME _mlirConversions
370  ADD_TO_PARENT MLIRPythonSources.Passes
371  ROOT_DIR "${PYTHON_SOURCE_DIR}"
372  SOURCES
373    Conversions/Conversions.cpp
374  PRIVATE_LINK_LIBS
375    LLVMSupport
376  EMBED_CAPI_LINK_LIBS
377  MLIRCAPIConversion
378)
379
380# Only enable the ExecutionEngine if the native target is configured in.
381if(TARGET ${LLVM_NATIVE_ARCH})
382  declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
383    MODULE_NAME _mlirExecutionEngine
384    ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
385    ROOT_DIR "${PYTHON_SOURCE_DIR}"
386    SOURCES
387      ExecutionEngineModule.cpp
388    PRIVATE_LINK_LIBS
389      LLVMSupport
390    EMBED_CAPI_LINK_LIBS
391      MLIRCAPIExecutionEngine
392  )
393endif()
394
395declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
396  MODULE_NAME _mlirGPUPasses
397  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
398  ROOT_DIR "${PYTHON_SOURCE_DIR}"
399  SOURCES
400    GPUPasses.cpp
401  PRIVATE_LINK_LIBS
402    LLVMSupport
403  EMBED_CAPI_LINK_LIBS
404    MLIRCAPIGPU
405)
406
407declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
408  MODULE_NAME _mlirLinalgPasses
409  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
410  ROOT_DIR "${PYTHON_SOURCE_DIR}"
411  SOURCES
412    LinalgPasses.cpp
413  PRIVATE_LINK_LIBS
414    LLVMSupport
415  EMBED_CAPI_LINK_LIBS
416    MLIRCAPILinalg
417)
418
419declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
420  MODULE_NAME _mlirSparseTensorPasses
421  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
422  ROOT_DIR "${PYTHON_SOURCE_DIR}"
423  SOURCES
424    SparseTensorPasses.cpp
425  PRIVATE_LINK_LIBS
426    LLVMSupport
427  EMBED_CAPI_LINK_LIBS
428    MLIRCAPISparseTensor
429)
430
431declare_mlir_python_extension(MLIRPythonExtension.Transforms
432  MODULE_NAME _mlirTransforms
433  ADD_TO_PARENT MLIRPythonSources.Passes
434  ROOT_DIR "${PYTHON_SOURCE_DIR}"
435  SOURCES
436    Transforms/Transforms.cpp
437  PRIVATE_LINK_LIBS
438    LLVMSupport
439  EMBED_CAPI_LINK_LIBS
440    MLIRCAPITransforms
441)
442
443# TODO: Figure out how to put this in the test tree.
444# This should not be included in the main Python extension. However,
445# putting it into MLIRPythonTestSources along with the dialect declaration
446# above confuses Python module loader when running under lit.
447set(_ADDL_TEST_SOURCES)
448if(MLIR_INCLUDE_TESTS)
449  set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
450  declare_mlir_python_sources(MLIRPythonTestSources)
451  declare_mlir_python_sources(MLIRPythonTestSources.Dialects
452    ADD_TO_PARENT MLIRPythonTestSources)
453
454  # TODO: this uses a tablegen file from the test directory and should be
455  # decoupled from here.
456  declare_mlir_python_sources(
457    MLIRPythonTestSources.Dialects.PythonTest
458    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
459    ADD_TO_PARENT MLIRPythonTestSources.Dialects
460    SOURCES dialects/python_test.py)
461  set(LLVM_TARGET_DEFINITIONS
462    "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
463  mlir_tablegen(
464    "dialects/_python_test_ops_gen.py"
465    -gen-python-op-bindings
466    -bind-dialect=python_test)
467  add_public_tablegen_target(PythonTestDialectPyIncGen)
468  declare_mlir_python_sources(
469    MLIRPythonTestSources.Dialects.PythonTest.ops_gen
470    ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
471    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
472    SOURCES "dialects/_python_test_ops_gen.py")
473
474  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
475    MODULE_NAME _mlirPythonTest
476    ADD_TO_PARENT MLIRPythonTestSources.Dialects
477    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
478    SOURCES
479      PythonTestModule.cpp
480    PRIVATE_LINK_LIBS
481      LLVMSupport
482    EMBED_CAPI_LINK_LIBS
483      MLIRCAPIPythonTestDialect
484  )
485endif()
486
487################################################################################
488# Common CAPI dependency DSO.
489# All python extensions must link through one DSO which exports the CAPI, and
490# this must have a globally unique name amongst all embeddors of the python
491# library since it will effectively have global scope.
492#
493# The presence of this aggregate library is part of the long term plan, but its
494# use needs to be made more flexible.
495#
496# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
497# once ready.
498################################################################################
499
500add_mlir_python_common_capi_library(MLIRPythonCAPI
501  INSTALL_COMPONENT MLIRPythonModules
502  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
503  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
504  RELATIVE_INSTALL_ROOT "../../../.."
505  DECLARED_HEADERS
506    MLIRPythonCAPI.HeaderSources
507  DECLARED_SOURCES
508    MLIRPythonSources
509    MLIRPythonExtension.AllPassesRegistration
510    ${_ADDL_TEST_SOURCES}
511)
512
513################################################################################
514# The fully assembled package of modules.
515# This must come last.
516################################################################################
517
518add_mlir_python_modules(MLIRPythonModules
519  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
520  INSTALL_PREFIX "python_packages/mlir_core/mlir"
521  DECLARED_SOURCES
522    MLIRPythonSources
523    MLIRPythonExtension.AllPassesRegistration
524    ${_ADDL_TEST_SOURCES}
525  COMMON_CAPI_LINK_LIBS
526    MLIRPythonCAPI
527)
528