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