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