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