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