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