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
11declare_mlir_python_sources(MLIRPythonTestSources)
12declare_mlir_python_sources(MLIRPythonTestSources.Dialects
13  ADD_TO_PARENT MLIRPythonTestSources)
14
15################################################################################
16# Pure python sources and generated code
17################################################################################
18
19declare_mlir_python_sources(MLIRPythonSources.Core
20  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
21  ADD_TO_PARENT MLIRPythonSources
22  SOURCES
23    _mlir_libs/__init__.py
24    ir.py
25    passmanager.py
26    dialects/_ods_common.py
27)
28
29declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
30  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
31  ADD_TO_PARENT MLIRPythonSources
32  SOURCES
33    execution_engine.py
34  SOURCES_GLOB
35    runtime/*.py
36)
37
38declare_mlir_python_sources(MLIRPythonSources.Passes
39  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
40  ADD_TO_PARENT MLIRPythonSources
41  SOURCES_GLOB
42    all_passes_registration/*.py
43    conversions/*.py
44    transforms/*.py
45)
46
47declare_mlir_python_sources(MLIRPythonCAPIHeaderSources
48  ROOT_DIR "${MLIR_SOURCE_DIR}/include"
49  SOURCES_GLOB "mlir-c/*.h"
50  DEST_PREFIX "_mlir_libs/include"
51)
52
53################################################################################
54# Dialect bindings
55################################################################################
56
57declare_mlir_dialect_python_bindings(
58  ADD_TO_PARENT MLIRPythonSources.Dialects
59  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
60  TD_FILE dialects/AsyncOps.td
61  SOURCES_GLOB dialects/async_dialect/*.py
62  DIALECT_NAME async_dialect)
63
64declare_mlir_dialect_python_bindings(
65  ADD_TO_PARENT MLIRPythonSources.Dialects
66  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
67  TD_FILE dialects/BuiltinOps.td
68  SOURCES
69    dialects/builtin.py
70    dialects/_builtin_ops_ext.py
71  DIALECT_NAME builtin)
72
73declare_mlir_dialect_python_bindings(
74  ADD_TO_PARENT MLIRPythonSources.Dialects
75  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
76  TD_FILE dialects/GPUOps.td
77  SOURCES_GLOB dialects/gpu/*.py
78  DIALECT_NAME gpu)
79
80declare_mlir_dialect_python_bindings(
81  ADD_TO_PARENT MLIRPythonSources.Dialects
82  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
83  TD_FILE dialects/LinalgOps.td
84  SOURCES
85    dialects/_linalg_ops_ext.py
86  SOURCES_GLOB
87    dialects/linalg/*.py
88  DIALECT_NAME linalg
89  DEPENDS LinalgOdsGen)
90
91declare_mlir_dialect_python_bindings(
92  ADD_TO_PARENT MLIRPythonSources.Dialects
93  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
94  TD_FILE dialects/MathOps.td
95  SOURCES dialects/math.py
96  DIALECT_NAME math)
97
98declare_mlir_dialect_python_bindings(
99  ADD_TO_PARENT MLIRPythonSources.Dialects
100  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
101  TD_FILE dialects/MemRefOps.td
102  SOURCES dialects/memref.py
103  DIALECT_NAME memref)
104
105declare_mlir_dialect_python_bindings(
106  ADD_TO_PARENT MLIRPythonTestSources.Dialects
107  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
108  TD_FILE dialects/PythonTest.td
109  SOURCES dialects/python_test.py
110  DIALECT_NAME python_test)
111
112declare_mlir_dialect_python_bindings(
113  ADD_TO_PARENT MLIRPythonSources.Dialects
114  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
115  TD_FILE dialects/SCFOps.td
116  SOURCES
117    dialects/scf.py
118    dialects/_scf_ops_ext.py
119  DIALECT_NAME scf)
120
121declare_mlir_dialect_python_bindings(
122  ADD_TO_PARENT MLIRPythonSources.Dialects
123  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
124  TD_FILE dialects/ShapeOps.td
125  SOURCES dialects/shape.py
126  DIALECT_NAME shape)
127
128declare_mlir_dialect_python_bindings(
129  ADD_TO_PARENT MLIRPythonSources.Dialects
130  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
131  TD_FILE dialects/SparseTensorOps.td
132  SOURCES dialects/sparse_tensor.py
133  DIALECT_NAME sparse_tensor)
134
135declare_mlir_dialect_python_bindings(
136  ADD_TO_PARENT MLIRPythonSources.Dialects
137  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
138  TD_FILE dialects/StandardOps.td
139  SOURCES dialects/std.py
140  DIALECT_NAME std)
141
142declare_mlir_dialect_python_bindings(
143  ADD_TO_PARENT MLIRPythonSources.Dialects
144  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
145  TD_FILE dialects/TensorOps.td
146  SOURCES dialects/tensor.py
147  DIALECT_NAME 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/TosaOps.td
153  SOURCES dialects/tosa.py
154  DIALECT_NAME tosa)
155
156declare_mlir_dialect_python_bindings(
157  ADD_TO_PARENT MLIRPythonSources.Dialects
158  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
159  TD_FILE dialects/VectorOps.td
160  SOURCES dialects/vector.py
161  DIALECT_NAME vector)
162
163################################################################################
164# Python extensions.
165# The sources for these are all in lib/Bindings/Python, but since they have to
166# be rebuilt for each package and integrate with the source setup here, we
167# just reference them here instead of having ordered, cross package target
168# dependencies.
169################################################################################
170
171set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
172declare_mlir_python_extension(MLIRPythonExtension.Core
173  MODULE_NAME _mlir
174  ADD_TO_PARENT MLIRPythonSources.Core
175  SOURCES
176    ${PYTHON_SOURCE_DIR}/DialectLinalg.cpp  # TODO: Break this out.
177    ${PYTHON_SOURCE_DIR}/DialectSparseTensor.cpp  # TODO: Break this out.
178    ${PYTHON_SOURCE_DIR}/MainModule.cpp
179    ${PYTHON_SOURCE_DIR}/IRAffine.cpp
180    ${PYTHON_SOURCE_DIR}/IRAttributes.cpp
181    ${PYTHON_SOURCE_DIR}/IRCore.cpp
182    ${PYTHON_SOURCE_DIR}/IRModule.cpp
183    ${PYTHON_SOURCE_DIR}/IRTypes.cpp
184    ${PYTHON_SOURCE_DIR}/PybindUtils.cpp
185    ${PYTHON_SOURCE_DIR}/Pass.cpp
186  PRIVATE_LINK_LIBS
187    LLVMSupport
188  EMBED_CAPI_LINK_LIBS
189    MLIRCAPIDebug
190    MLIRCAPIIR
191    MLIRCAPIRegistration  # TODO: See about dis-aggregating
192
193    # Dialects
194    MLIRCAPILinalg  # TODO: Remove when above is removed.
195    MLIRCAPISparseTensor  # TODO: Remove when above is removed.
196    MLIRCAPIStandard
197)
198
199declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
200  MODULE_NAME _mlirAllPassesRegistration
201  SOURCES
202    ${PYTHON_SOURCE_DIR}/AllPassesRegistration.cpp
203  PRIVATE_LINK_LIBS
204    LLVMSupport
205  EMBED_CAPI_LINK_LIBS
206    MLIRCAPIConversion
207    MLIRCAPITransforms
208)
209
210declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
211  MODULE_NAME _mlirAsyncPasses
212  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
213  SOURCES
214    ${PYTHON_SOURCE_DIR}/AsyncPasses.cpp
215  PRIVATE_LINK_LIBS
216    LLVMSupport
217  EMBED_CAPI_LINK_LIBS
218    MLIRCAPIAsync
219)
220
221declare_mlir_python_extension(MLIRPythonExtension.Conversions
222  MODULE_NAME _mlirConversions
223  ADD_TO_PARENT MLIRPythonSources.Passes
224  SOURCES
225    ${PYTHON_SOURCE_DIR}/Conversions/Conversions.cpp
226  PRIVATE_LINK_LIBS
227    LLVMSupport
228  EMBED_CAPI_LINK_LIBS
229  MLIRCAPIConversion
230)
231
232declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
233  MODULE_NAME _mlirExecutionEngine
234  ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
235  SOURCES
236    ${PYTHON_SOURCE_DIR}/ExecutionEngineModule.cpp
237  PRIVATE_LINK_LIBS
238    LLVMSupport
239  EMBED_CAPI_LINK_LIBS
240    MLIRCEXECUTIONENGINE
241)
242
243declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
244  MODULE_NAME _mlirGPUPasses
245  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
246  SOURCES
247    ${PYTHON_SOURCE_DIR}/GPUPasses.cpp
248  PRIVATE_LINK_LIBS
249    LLVMSupport
250  EMBED_CAPI_LINK_LIBS
251    MLIRCAPIGPU
252)
253
254declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
255  MODULE_NAME _mlirLinalgPasses
256  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
257  SOURCES
258    ${PYTHON_SOURCE_DIR}/LinalgPasses.cpp
259  PRIVATE_LINK_LIBS
260    LLVMSupport
261  EMBED_CAPI_LINK_LIBS
262    MLIRCAPILinalg
263)
264
265declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
266  MODULE_NAME _mlirSparseTensorPasses
267  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
268  SOURCES
269    ${PYTHON_SOURCE_DIR}/SparseTensorPasses.cpp
270  PRIVATE_LINK_LIBS
271    LLVMSupport
272  EMBED_CAPI_LINK_LIBS
273    MLIRCAPISparseTensor
274)
275
276declare_mlir_python_extension(MLIRPythonExtension.Transforms
277  MODULE_NAME _mlirTransforms
278  ADD_TO_PARENT MLIRPythonSources.Passes
279  SOURCES
280    ${PYTHON_SOURCE_DIR}/Transforms/Transforms.cpp
281  PRIVATE_LINK_LIBS
282    LLVMSupport
283  EMBED_CAPI_LINK_LIBS
284    MLIRCAPITransforms
285)
286
287################################################################################
288# Common CAPI dependency DSO.
289# All python extensions must link through one DSO which exports the CAPI, and
290# this must have a globally unique name amongst all embeddors of the python
291# library since it will effectively have global scope.
292#
293# The presence of this aggregate library is part of the long term plan, but its
294# use needs to be made more flexible.
295#
296# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
297# once ready.
298################################################################################
299
300add_mlir_python_common_capi_library(MLIRPythonCAPI
301  INSTALL_COMPONENT MLIRPythonModules
302  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
303  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
304  RELATIVE_INSTALL_ROOT "../../../.."
305  DECLARED_SOURCES
306    MLIRPythonSources
307    MLIRPythonExtension.AllPassesRegistration
308)
309
310################################################################################
311# The fully assembled package of modules.
312# This must come last.
313################################################################################
314
315add_mlir_python_modules(MLIRPythonModules
316  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
317  INSTALL_PREFIX "python_packages/mlir_core/mlir"
318  DECLARED_SOURCES
319    MLIRPythonSources
320    MLIRPythonExtension.AllPassesRegistration
321    MLIRPythonCAPIHeaderSources
322  COMMON_CAPI_LINK_LIBS
323    MLIRPythonCAPI
324  )
325
326
327add_mlir_python_modules(MLIRPythonTestModules
328  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_test/mlir"
329  INSTALL_PREFIX "python_packages/mlir_test/mlir"
330  DECLARED_SOURCES
331    MLIRPythonTestSources
332  )
333