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