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