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
11################################################################################
12# Pure python sources and generated code
13################################################################################
14
15declare_mlir_python_sources(MLIRPythonSources.Core
16  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
17  ADD_TO_PARENT MLIRPythonSources
18  SOURCES
19    _mlir_libs/__init__.py
20    ir.py
21    passmanager.py
22    dialects/_ods_common.py
23
24    # The main _mlir module has submodules: include stubs from each.
25    _mlir_libs/_mlir/__init__.pyi
26    _mlir_libs/_mlir/ir.pyi
27    _mlir_libs/_mlir/passmanager.pyi
28)
29
30declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
31  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
32  ADD_TO_PARENT MLIRPythonSources
33  SOURCES
34    execution_engine.py
35    _mlir_libs/_mlirExecutionEngine.pyi
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/ArithmeticOps.td
104  SOURCES
105    dialects/arith.py
106    dialects/_arith_ops_ext.py
107  DIALECT_NAME arith)
108
109declare_mlir_dialect_python_bindings(
110  ADD_TO_PARENT MLIRPythonSources.Dialects
111  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
112  TD_FILE dialects/MemRefOps.td
113  SOURCES
114    dialects/memref.py
115    dialects/_memref_ops_ext.py
116  DIALECT_NAME memref)
117
118declare_mlir_python_sources(
119  MLIRPythonSources.Dialects.quant
120  ADD_TO_PARENT MLIRPythonSources.Dialects
121  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
122  SOURCES
123    dialects/quant.py
124    _mlir_libs/_mlir/dialects/quant.pyi)
125
126declare_mlir_dialect_python_bindings(
127  ADD_TO_PARENT MLIRPythonSources.Dialects
128  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
129  TD_FILE dialects/SCFOps.td
130  SOURCES
131    dialects/scf.py
132    dialects/_scf_ops_ext.py
133  DIALECT_NAME scf)
134
135declare_mlir_dialect_python_bindings(
136  ADD_TO_PARENT MLIRPythonSources.Dialects
137  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
138  TD_FILE dialects/ShapeOps.td
139  SOURCES dialects/shape.py
140  DIALECT_NAME shape)
141
142declare_mlir_dialect_python_bindings(
143  ADD_TO_PARENT MLIRPythonSources.Dialects
144  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
145  TD_FILE dialects/SparseTensorOps.td
146  SOURCES dialects/sparse_tensor.py
147  DIALECT_NAME sparse_tensor)
148
149declare_mlir_dialect_python_bindings(
150  ADD_TO_PARENT MLIRPythonSources.Dialects
151  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
152  TD_FILE dialects/StandardOps.td
153  SOURCES
154    dialects/std.py
155    dialects/_std_ops_ext.py
156  DIALECT_NAME std)
157
158declare_mlir_dialect_python_bindings(
159  ADD_TO_PARENT MLIRPythonSources.Dialects
160  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
161  TD_FILE dialects/TensorOps.td
162  SOURCES dialects/tensor.py
163  DIALECT_NAME tensor)
164
165declare_mlir_dialect_python_bindings(
166  ADD_TO_PARENT MLIRPythonSources.Dialects
167  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
168  TD_FILE dialects/TosaOps.td
169  SOURCES dialects/tosa.py
170  DIALECT_NAME tosa)
171
172declare_mlir_dialect_python_bindings(
173  ADD_TO_PARENT MLIRPythonSources.Dialects
174  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
175  TD_FILE dialects/VectorOps.td
176  SOURCES dialects/vector.py
177  DIALECT_NAME vector)
178
179################################################################################
180# Python extensions.
181# The sources for these are all in lib/Bindings/Python, but since they have to
182# be rebuilt for each package and integrate with the source setup here, we
183# just reference them here instead of having ordered, cross package target
184# dependencies.
185################################################################################
186
187set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
188declare_mlir_python_extension(MLIRPythonExtension.Core
189  MODULE_NAME _mlir
190  ADD_TO_PARENT MLIRPythonSources.Core
191  ROOT_DIR "${PYTHON_SOURCE_DIR}"
192  SOURCES
193    MainModule.cpp
194    IRAffine.cpp
195    IRAttributes.cpp
196    IRCore.cpp
197    IRInterfaces.cpp
198    IRModule.cpp
199    IRTypes.cpp
200    PybindUtils.cpp
201    Pass.cpp
202
203    # Headers must be included explicitly so they are installed.
204    Globals.h
205    IRModule.h
206    Pass.h
207    PybindUtils.h
208  PRIVATE_LINK_LIBS
209    LLVMSupport
210  EMBED_CAPI_LINK_LIBS
211    MLIRCAPIDebug
212    MLIRCAPIIR
213    MLIRCAPIInterfaces
214    MLIRCAPIRegistration  # TODO: See about dis-aggregating
215
216    # Dialects
217    MLIRCAPIStandard
218)
219
220declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
221  MODULE_NAME _mlirDialectsLinalg
222  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
223  ROOT_DIR "${PYTHON_SOURCE_DIR}"
224  SOURCES
225    DialectLinalg.cpp
226  PRIVATE_LINK_LIBS
227    LLVMSupport
228  EMBED_CAPI_LINK_LIBS
229    MLIRCAPIIR
230    MLIRCAPILinalg
231)
232
233declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
234  MODULE_NAME _mlirDialectsQuant
235  ADD_TO_PARENT MLIRPythonSources.Dialects.quant
236  ROOT_DIR "${PYTHON_SOURCE_DIR}"
237  SOURCES
238    DialectQuant.cpp
239  PRIVATE_LINK_LIBS
240    LLVMSupport
241  EMBED_CAPI_LINK_LIBS
242    MLIRCAPIIR
243    MLIRCAPIQuant
244)
245
246declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
247  MODULE_NAME _mlirDialectsSparseTensor
248  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
249  ROOT_DIR "${PYTHON_SOURCE_DIR}"
250  SOURCES
251    DialectSparseTensor.cpp
252  PRIVATE_LINK_LIBS
253    LLVMSupport
254  EMBED_CAPI_LINK_LIBS
255    MLIRCAPIIR
256    MLIRCAPISparseTensor
257)
258
259declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
260  MODULE_NAME _mlirAllPassesRegistration
261  ROOT_DIR "${PYTHON_SOURCE_DIR}"
262  SOURCES
263    AllPassesRegistration.cpp
264  PRIVATE_LINK_LIBS
265    LLVMSupport
266  EMBED_CAPI_LINK_LIBS
267    MLIRCAPIConversion
268    MLIRCAPITransforms
269)
270
271declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
272  MODULE_NAME _mlirAsyncPasses
273  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
274  ROOT_DIR "${PYTHON_SOURCE_DIR}"
275  SOURCES
276    AsyncPasses.cpp
277  PRIVATE_LINK_LIBS
278    LLVMSupport
279  EMBED_CAPI_LINK_LIBS
280    MLIRCAPIAsync
281)
282
283declare_mlir_python_extension(MLIRPythonExtension.Conversions
284  MODULE_NAME _mlirConversions
285  ADD_TO_PARENT MLIRPythonSources.Passes
286  ROOT_DIR "${PYTHON_SOURCE_DIR}"
287  SOURCES
288    Conversions/Conversions.cpp
289  PRIVATE_LINK_LIBS
290    LLVMSupport
291  EMBED_CAPI_LINK_LIBS
292  MLIRCAPIConversion
293)
294
295# Only enable the ExecutionEngine if the native target is configured in.
296if(TARGET ${LLVM_NATIVE_ARCH})
297  declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
298    MODULE_NAME _mlirExecutionEngine
299    ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
300    ROOT_DIR "${PYTHON_SOURCE_DIR}"
301    SOURCES
302      ExecutionEngineModule.cpp
303    PRIVATE_LINK_LIBS
304      LLVMSupport
305    EMBED_CAPI_LINK_LIBS
306      MLIRCAPIExecutionEngine
307  )
308endif()
309
310declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
311  MODULE_NAME _mlirGPUPasses
312  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
313  ROOT_DIR "${PYTHON_SOURCE_DIR}"
314  SOURCES
315    GPUPasses.cpp
316  PRIVATE_LINK_LIBS
317    LLVMSupport
318  EMBED_CAPI_LINK_LIBS
319    MLIRCAPIGPU
320)
321
322declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
323  MODULE_NAME _mlirLinalgPasses
324  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
325  ROOT_DIR "${PYTHON_SOURCE_DIR}"
326  SOURCES
327    LinalgPasses.cpp
328  PRIVATE_LINK_LIBS
329    LLVMSupport
330  EMBED_CAPI_LINK_LIBS
331    MLIRCAPILinalg
332)
333
334declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
335  MODULE_NAME _mlirSparseTensorPasses
336  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
337  ROOT_DIR "${PYTHON_SOURCE_DIR}"
338  SOURCES
339    SparseTensorPasses.cpp
340  PRIVATE_LINK_LIBS
341    LLVMSupport
342  EMBED_CAPI_LINK_LIBS
343    MLIRCAPISparseTensor
344)
345
346declare_mlir_python_extension(MLIRPythonExtension.Transforms
347  MODULE_NAME _mlirTransforms
348  ADD_TO_PARENT MLIRPythonSources.Passes
349  ROOT_DIR "${PYTHON_SOURCE_DIR}"
350  SOURCES
351    Transforms/Transforms.cpp
352  PRIVATE_LINK_LIBS
353    LLVMSupport
354  EMBED_CAPI_LINK_LIBS
355    MLIRCAPITransforms
356)
357
358# TODO: Figure out how to put this in the test tree.
359# This should not be included in the main Python extension. However,
360# putting it into MLIRPythonTestSources along with the dialect declaration
361# above confuses Python module loader when running under lit.
362set(_ADDL_TEST_SOURCES)
363if(MLIR_INCLUDE_TESTS)
364  set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
365  declare_mlir_python_sources(MLIRPythonTestSources)
366  declare_mlir_python_sources(MLIRPythonTestSources.Dialects
367    ADD_TO_PARENT MLIRPythonTestSources)
368
369  # TODO: this uses a tablegen file from the test directory and should be
370  # decoupled from here.
371  declare_mlir_python_sources(
372    MLIRPythonTestSources.Dialects.PythonTest
373    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
374    ADD_TO_PARENT MLIRPythonTestSources.Dialects
375    SOURCES dialects/python_test.py)
376  set(LLVM_TARGET_DEFINITIONS
377    "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
378  mlir_tablegen(
379    "dialects/_python_test_ops_gen.py"
380    -gen-python-op-bindings
381    -bind-dialect=python_test)
382  add_public_tablegen_target(PythonTestDialectPyIncGen)
383  declare_mlir_python_sources(
384    MLIRPythonTestSources.Dialects.PythonTest.ops_gen
385    ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
386    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
387    SOURCES "dialects/_python_test_ops_gen.py")
388
389  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
390    MODULE_NAME _mlirPythonTest
391    ADD_TO_PARENT MLIRPythonTestSources.Dialects
392    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
393    SOURCES
394      PythonTestModule.cpp
395    PRIVATE_LINK_LIBS
396      LLVMSupport
397    EMBED_CAPI_LINK_LIBS
398      MLIRCAPIPythonTestDialect
399  )
400endif()
401
402################################################################################
403# Common CAPI dependency DSO.
404# All python extensions must link through one DSO which exports the CAPI, and
405# this must have a globally unique name amongst all embeddors of the python
406# library since it will effectively have global scope.
407#
408# The presence of this aggregate library is part of the long term plan, but its
409# use needs to be made more flexible.
410#
411# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
412# once ready.
413################################################################################
414
415add_mlir_python_common_capi_library(MLIRPythonCAPI
416  INSTALL_COMPONENT MLIRPythonModules
417  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
418  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
419  RELATIVE_INSTALL_ROOT "../../../.."
420  DECLARED_SOURCES
421    MLIRPythonSources
422    MLIRPythonExtension.AllPassesRegistration
423    ${_ADDL_TEST_SOURCES}
424)
425
426################################################################################
427# The fully assembled package of modules.
428# This must come last.
429################################################################################
430
431add_mlir_python_modules(MLIRPythonModules
432  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
433  INSTALL_PREFIX "python_packages/mlir_core/mlir"
434  DECLARED_SOURCES
435    MLIRPythonSources
436    MLIRPythonExtension.AllPassesRegistration
437    MLIRPythonCAPIHeaderSources
438    ${_ADDL_TEST_SOURCES}
439  COMMON_CAPI_LINK_LIBS
440    MLIRPythonCAPI
441  )
442