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