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