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/ShapeOps.td
116  SOURCES dialects/shape.py
117  DIALECT_NAME shape)
118
119declare_mlir_dialect_python_bindings(
120  ADD_TO_PARENT MLIRPythonSources.Dialects
121  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
122  SOURCES dialects/sparse_tensor.py
123  DIALECT_NAME sparse_tensor)
124
125declare_mlir_dialect_python_bindings(
126  ADD_TO_PARENT MLIRPythonSources.Dialects
127  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
128  TD_FILE dialects/StandardOps.td
129  SOURCES dialects/std.py
130  DIALECT_NAME std)
131
132declare_mlir_dialect_python_bindings(
133  ADD_TO_PARENT MLIRPythonSources.Dialects
134  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
135  TD_FILE dialects/TensorOps.td
136  SOURCES dialects/tensor.py
137  DIALECT_NAME tensor)
138
139declare_mlir_dialect_python_bindings(
140  ADD_TO_PARENT MLIRPythonSources.Dialects
141  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
142  TD_FILE dialects/TosaOps.td
143  SOURCES dialects/tosa.py
144  DIALECT_NAME tosa)
145
146declare_mlir_dialect_python_bindings(
147  ADD_TO_PARENT MLIRPythonSources.Dialects
148  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
149  TD_FILE dialects/VectorOps.td
150  SOURCES dialects/vector.py
151  DIALECT_NAME vector)
152
153################################################################################
154# Python extensions.
155# The sources for these are all in lib/Bindings/Python, but since they have to
156# be rebuilt for each package and integrate with the source setup here, we
157# just reference them here instead of having ordered, cross package target
158# dependencies.
159################################################################################
160
161set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
162declare_mlir_python_extension(MLIRPythonExtension.Core
163  MODULE_NAME _mlir
164  ADD_TO_PARENT MLIRPythonSources.Core
165  SOURCES
166    ${PYTHON_SOURCE_DIR}/DialectLinalg.cpp  # TODO: Break this out.
167    ${PYTHON_SOURCE_DIR}/DialectSparseTensor.cpp  # TODO: Break this out.
168    ${PYTHON_SOURCE_DIR}/MainModule.cpp
169    ${PYTHON_SOURCE_DIR}/IRAffine.cpp
170    ${PYTHON_SOURCE_DIR}/IRAttributes.cpp
171    ${PYTHON_SOURCE_DIR}/IRCore.cpp
172    ${PYTHON_SOURCE_DIR}/IRModule.cpp
173    ${PYTHON_SOURCE_DIR}/IRTypes.cpp
174    ${PYTHON_SOURCE_DIR}/PybindUtils.cpp
175    ${PYTHON_SOURCE_DIR}/Pass.cpp
176  PRIVATE_LINK_LIBS
177    LLVMSupport
178  EMBED_CAPI_LINK_LIBS
179    MLIRCAPIDebug
180    MLIRCAPIIR
181    MLIRCAPIRegistration  # TODO: See about dis-aggregating
182
183    # Dialects
184    MLIRCAPILinalg  # TODO: Remove when above is removed.
185    MLIRCAPISparseTensor  # TODO: Remove when above is removed.
186    MLIRCAPIStandard
187)
188
189declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
190  MODULE_NAME _mlirAllPassesRegistration
191  SOURCES
192    ${PYTHON_SOURCE_DIR}/AllPassesRegistration.cpp
193  PRIVATE_LINK_LIBS
194    LLVMSupport
195  EMBED_CAPI_LINK_LIBS
196    MLIRCAPIConversion
197    MLIRCAPITransforms
198)
199
200declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
201  MODULE_NAME _mlirAsyncPasses
202  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
203  SOURCES
204    ${PYTHON_SOURCE_DIR}/AsyncPasses.cpp
205  PRIVATE_LINK_LIBS
206    LLVMSupport
207  EMBED_CAPI_LINK_LIBS
208    MLIRCAPIAsync
209)
210
211declare_mlir_python_extension(MLIRPythonExtension.Conversions
212  MODULE_NAME _mlirConversions
213  ADD_TO_PARENT MLIRPythonSources.Passes
214  SOURCES
215    ${PYTHON_SOURCE_DIR}/Conversions/Conversions.cpp
216  PRIVATE_LINK_LIBS
217    LLVMSupport
218  EMBED_CAPI_LINK_LIBS
219  MLIRCAPIConversion
220)
221
222declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
223  MODULE_NAME _mlirExecutionEngine
224  ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
225  SOURCES
226    ${PYTHON_SOURCE_DIR}/ExecutionEngineModule.cpp
227  PRIVATE_LINK_LIBS
228    LLVMSupport
229  EMBED_CAPI_LINK_LIBS
230    MLIRCEXECUTIONENGINE
231)
232
233declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
234  MODULE_NAME _mlirGPUPasses
235  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
236  SOURCES
237    ${PYTHON_SOURCE_DIR}/GPUPasses.cpp
238  PRIVATE_LINK_LIBS
239    LLVMSupport
240  EMBED_CAPI_LINK_LIBS
241    MLIRCAPIGPU
242)
243
244declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
245  MODULE_NAME _mlirLinalgPasses
246  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
247  SOURCES
248    ${PYTHON_SOURCE_DIR}/LinalgPasses.cpp
249  PRIVATE_LINK_LIBS
250    LLVMSupport
251  EMBED_CAPI_LINK_LIBS
252    MLIRCAPILinalg
253)
254
255declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
256  MODULE_NAME _mlirSparseTensorPasses
257  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
258  SOURCES
259    ${PYTHON_SOURCE_DIR}/SparseTensorPasses.cpp
260  PRIVATE_LINK_LIBS
261    LLVMSupport
262  EMBED_CAPI_LINK_LIBS
263    MLIRCAPISparseTensor
264)
265
266declare_mlir_python_extension(MLIRPythonExtension.Transforms
267  MODULE_NAME _mlirTransforms
268  ADD_TO_PARENT MLIRPythonSources.Passes
269  SOURCES
270    ${PYTHON_SOURCE_DIR}/Transforms/Transforms.cpp
271  PRIVATE_LINK_LIBS
272    LLVMSupport
273  EMBED_CAPI_LINK_LIBS
274    MLIRCAPITransforms
275)
276
277################################################################################
278# Common CAPI dependency DSO.
279# All python extensions must link through one DSO which exports the CAPI, and
280# this must have a globally unique name amongst all embeddors of the python
281# library since it will effectively have global scope.
282#
283# The presence of this aggregate library is part of the long term plan, but its
284# use needs to be made more flexible.
285#
286# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
287# once ready.
288################################################################################
289
290add_mlir_python_common_capi_library(MLIRPythonCAPI
291  INSTALL_COMPONENT MLIRPythonModules
292  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
293  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
294  RELATIVE_INSTALL_ROOT "../../../.."
295  DECLARED_SOURCES
296    MLIRPythonSources
297    MLIRPythonExtension.AllPassesRegistration
298)
299
300################################################################################
301# The fully assembled package of modules.
302# This must come last.
303################################################################################
304
305add_mlir_python_modules(MLIRPythonModules
306  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
307  INSTALL_PREFIX "python_packages/mlir_core/mlir"
308  DECLARED_SOURCES
309    MLIRPythonSources
310    MLIRPythonExtension.AllPassesRegistration
311    MLIRPythonCAPIHeaderSources
312  COMMON_CAPI_LINK_LIBS
313    MLIRPythonCAPI
314  )
315
316
317add_mlir_python_modules(MLIRPythonTestModules
318  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_test/mlir"
319  INSTALL_PREFIX "python_packages/mlir_test/mlir"
320  DECLARED_SOURCES
321    MLIRPythonTestSources
322  )
323