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