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