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
366# Only enable the ExecutionEngine if the native target is configured in.
367if(TARGET ${LLVM_NATIVE_ARCH})
368  declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
369    MODULE_NAME _mlirExecutionEngine
370    ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
371    ROOT_DIR "${PYTHON_SOURCE_DIR}"
372    SOURCES
373      ExecutionEngineModule.cpp
374    PRIVATE_LINK_LIBS
375      LLVMSupport
376    EMBED_CAPI_LINK_LIBS
377      MLIRCAPIExecutionEngine
378  )
379endif()
380
381declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
382  MODULE_NAME _mlirGPUPasses
383  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
384  ROOT_DIR "${PYTHON_SOURCE_DIR}"
385  SOURCES
386    GPUPasses.cpp
387  PRIVATE_LINK_LIBS
388    LLVMSupport
389  EMBED_CAPI_LINK_LIBS
390    MLIRCAPIGPU
391)
392
393declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
394  MODULE_NAME _mlirLinalgPasses
395  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
396  ROOT_DIR "${PYTHON_SOURCE_DIR}"
397  SOURCES
398    LinalgPasses.cpp
399  PRIVATE_LINK_LIBS
400    LLVMSupport
401  EMBED_CAPI_LINK_LIBS
402    MLIRCAPILinalg
403)
404
405declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
406  MODULE_NAME _mlirSparseTensorPasses
407  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
408  ROOT_DIR "${PYTHON_SOURCE_DIR}"
409  SOURCES
410    SparseTensorPasses.cpp
411  PRIVATE_LINK_LIBS
412    LLVMSupport
413  EMBED_CAPI_LINK_LIBS
414    MLIRCAPISparseTensor
415)
416
417# TODO: Figure out how to put this in the test tree.
418# This should not be included in the main Python extension. However,
419# putting it into MLIRPythonTestSources along with the dialect declaration
420# above confuses Python module loader when running under lit.
421set(_ADDL_TEST_SOURCES)
422if(MLIR_INCLUDE_TESTS)
423  set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
424  declare_mlir_python_sources(MLIRPythonTestSources)
425  declare_mlir_python_sources(MLIRPythonTestSources.Dialects
426    ADD_TO_PARENT MLIRPythonTestSources)
427
428  # TODO: this uses a tablegen file from the test directory and should be
429  # decoupled from here.
430  declare_mlir_python_sources(
431    MLIRPythonTestSources.Dialects.PythonTest
432    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
433    ADD_TO_PARENT MLIRPythonTestSources.Dialects
434    SOURCES dialects/python_test.py)
435  set(LLVM_TARGET_DEFINITIONS
436    "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
437  mlir_tablegen(
438    "dialects/_python_test_ops_gen.py"
439    -gen-python-op-bindings
440    -bind-dialect=python_test)
441  add_public_tablegen_target(PythonTestDialectPyIncGen)
442  declare_mlir_python_sources(
443    MLIRPythonTestSources.Dialects.PythonTest.ops_gen
444    ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
445    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
446    SOURCES "dialects/_python_test_ops_gen.py")
447
448  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
449    MODULE_NAME _mlirPythonTest
450    ADD_TO_PARENT MLIRPythonTestSources.Dialects
451    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
452    SOURCES
453      PythonTestModule.cpp
454    PRIVATE_LINK_LIBS
455      LLVMSupport
456    EMBED_CAPI_LINK_LIBS
457      MLIRCAPIPythonTestDialect
458  )
459endif()
460
461################################################################################
462# Common CAPI dependency DSO.
463# All python extensions must link through one DSO which exports the CAPI, and
464# this must have a globally unique name amongst all embeddors of the python
465# library since it will effectively have global scope.
466#
467# The presence of this aggregate library is part of the long term plan, but its
468# use needs to be made more flexible.
469#
470# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
471# once ready.
472################################################################################
473
474add_mlir_python_common_capi_library(MLIRPythonCAPI
475  INSTALL_COMPONENT MLIRPythonModules
476  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
477  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
478  RELATIVE_INSTALL_ROOT "../../../.."
479  DECLARED_HEADERS
480    MLIRPythonCAPI.HeaderSources
481  DECLARED_SOURCES
482    MLIRPythonSources
483    MLIRPythonExtension.RegisterEverything
484    ${_ADDL_TEST_SOURCES}
485)
486
487################################################################################
488# The fully assembled package of modules.
489# This must come last.
490################################################################################
491
492add_mlir_python_modules(MLIRPythonModules
493  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
494  INSTALL_PREFIX "python_packages/mlir_core/mlir"
495  DECLARED_SOURCES
496    MLIRPythonSources
497    MLIRPythonExtension.RegisterEverything
498    ${_ADDL_TEST_SOURCES}
499  COMMON_CAPI_LINK_LIBS
500    MLIRPythonCAPI
501)
502