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