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