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
25declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
26  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
27  ADD_TO_PARENT MLIRPythonSources
28  SOURCES
29    execution_engine.py
30  SOURCES_GLOB
31    runtime/*.py
32)
33
34declare_mlir_python_sources(MLIRPythonSources.Passes
35  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
36  ADD_TO_PARENT MLIRPythonSources
37  SOURCES_GLOB
38    all_passes_registration/*.py
39    conversions/*.py
40    transforms/*.py
41)
42
43declare_mlir_python_sources(MLIRPythonCAPIHeaderSources
44  ROOT_DIR "${MLIR_SOURCE_DIR}/include"
45  SOURCES_GLOB "mlir-c/*.h"
46  DEST_PREFIX "_mlir_libs/include"
47)
48
49################################################################################
50# Dialect bindings
51################################################################################
52
53declare_mlir_dialect_python_bindings(
54  ADD_TO_PARENT MLIRPythonSources.Dialects
55  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
56  TD_FILE dialects/AsyncOps.td
57  SOURCES_GLOB dialects/async_dialect/*.py
58  DIALECT_NAME async_dialect)
59
60declare_mlir_dialect_python_bindings(
61  ADD_TO_PARENT MLIRPythonSources.Dialects
62  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
63  TD_FILE dialects/BuiltinOps.td
64  SOURCES
65    dialects/builtin.py
66    dialects/_builtin_ops_ext.py
67  DIALECT_NAME builtin)
68
69declare_mlir_dialect_python_bindings(
70  ADD_TO_PARENT MLIRPythonSources.Dialects
71  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
72  TD_FILE dialects/GPUOps.td
73  SOURCES_GLOB dialects/gpu/*.py
74  DIALECT_NAME gpu)
75
76declare_mlir_dialect_python_bindings(
77  ADD_TO_PARENT MLIRPythonSources.Dialects
78  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
79  TD_FILE dialects/LinalgOps.td
80  SOURCES
81    dialects/_linalg_ops_ext.py
82  SOURCES_GLOB
83    dialects/linalg/*.py
84  DIALECT_NAME linalg
85  DEPENDS LinalgOdsGen)
86
87declare_mlir_dialect_python_bindings(
88  ADD_TO_PARENT MLIRPythonSources.Dialects
89  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
90  TD_FILE dialects/MathOps.td
91  SOURCES dialects/math.py
92  DIALECT_NAME math)
93
94declare_mlir_dialect_python_bindings(
95  ADD_TO_PARENT MLIRPythonSources.Dialects
96  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
97  TD_FILE dialects/ArithmeticOps.td
98  SOURCES
99    dialects/arith.py
100    dialects/_arith_ops_ext.py
101  DIALECT_NAME arith)
102
103declare_mlir_dialect_python_bindings(
104  ADD_TO_PARENT MLIRPythonSources.Dialects
105  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
106  TD_FILE dialects/MemRefOps.td
107  SOURCES
108    dialects/memref.py
109    dialects/_memref_ops_ext.py
110  DIALECT_NAME memref)
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  ROOT_DIR "${PYTHON_SOURCE_DIR}"
178  SOURCES
179    DialectLinalg.cpp  # TODO: Break this out.
180    DialectSparseTensor.cpp  # TODO: Break this out.
181    MainModule.cpp
182    IRAffine.cpp
183    IRAttributes.cpp
184    IRCore.cpp
185    IRInterfaces.cpp
186    IRModule.cpp
187    IRTypes.cpp
188    PybindUtils.cpp
189    Pass.cpp
190
191    # Headers must be included explicitly so they are installed.
192    Dialects.h
193    Globals.h
194    IRModule.h
195    Pass.h
196    PybindUtils.h
197  PRIVATE_LINK_LIBS
198    LLVMSupport
199  EMBED_CAPI_LINK_LIBS
200    MLIRCAPIDebug
201    MLIRCAPIIR
202    MLIRCAPIInterfaces
203    MLIRCAPIRegistration  # TODO: See about dis-aggregating
204
205    # Dialects
206    MLIRCAPILinalg  # TODO: Remove when above is removed.
207    MLIRCAPISparseTensor  # TODO: Remove when above is removed.
208    MLIRCAPIStandard
209)
210
211declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
212  MODULE_NAME _mlirAllPassesRegistration
213  ROOT_DIR "${PYTHON_SOURCE_DIR}"
214  SOURCES
215    AllPassesRegistration.cpp
216  PRIVATE_LINK_LIBS
217    LLVMSupport
218  EMBED_CAPI_LINK_LIBS
219    MLIRCAPIConversion
220    MLIRCAPITransforms
221)
222
223declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
224  MODULE_NAME _mlirAsyncPasses
225  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
226  ROOT_DIR "${PYTHON_SOURCE_DIR}"
227  SOURCES
228    AsyncPasses.cpp
229  PRIVATE_LINK_LIBS
230    LLVMSupport
231  EMBED_CAPI_LINK_LIBS
232    MLIRCAPIAsync
233)
234
235declare_mlir_python_extension(MLIRPythonExtension.Conversions
236  MODULE_NAME _mlirConversions
237  ADD_TO_PARENT MLIRPythonSources.Passes
238  ROOT_DIR "${PYTHON_SOURCE_DIR}"
239  SOURCES
240    Conversions/Conversions.cpp
241  PRIVATE_LINK_LIBS
242    LLVMSupport
243  EMBED_CAPI_LINK_LIBS
244  MLIRCAPIConversion
245)
246
247declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
248  MODULE_NAME _mlirExecutionEngine
249  ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
250  ROOT_DIR "${PYTHON_SOURCE_DIR}"
251  SOURCES
252    ExecutionEngineModule.cpp
253  PRIVATE_LINK_LIBS
254    LLVMSupport
255  EMBED_CAPI_LINK_LIBS
256    MLIRCEXECUTIONENGINE
257)
258
259declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
260  MODULE_NAME _mlirGPUPasses
261  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
262  ROOT_DIR "${PYTHON_SOURCE_DIR}"
263  SOURCES
264    GPUPasses.cpp
265  PRIVATE_LINK_LIBS
266    LLVMSupport
267  EMBED_CAPI_LINK_LIBS
268    MLIRCAPIGPU
269)
270
271declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
272  MODULE_NAME _mlirLinalgPasses
273  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
274  ROOT_DIR "${PYTHON_SOURCE_DIR}"
275  SOURCES
276    LinalgPasses.cpp
277  PRIVATE_LINK_LIBS
278    LLVMSupport
279  EMBED_CAPI_LINK_LIBS
280    MLIRCAPILinalg
281)
282
283declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
284  MODULE_NAME _mlirSparseTensorPasses
285  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
286  ROOT_DIR "${PYTHON_SOURCE_DIR}"
287  SOURCES
288    SparseTensorPasses.cpp
289  PRIVATE_LINK_LIBS
290    LLVMSupport
291  EMBED_CAPI_LINK_LIBS
292    MLIRCAPISparseTensor
293)
294
295declare_mlir_python_extension(MLIRPythonExtension.Transforms
296  MODULE_NAME _mlirTransforms
297  ADD_TO_PARENT MLIRPythonSources.Passes
298  ROOT_DIR "${PYTHON_SOURCE_DIR}"
299  SOURCES
300    Transforms/Transforms.cpp
301  PRIVATE_LINK_LIBS
302    LLVMSupport
303  EMBED_CAPI_LINK_LIBS
304    MLIRCAPITransforms
305)
306
307# TODO: Figure out how to put this in the test tree.
308# This should not be included in the main Python extension. However,
309# putting it into MLIRPythonTestSources along with the dialect declaration
310# above confuses Python module loader when running under lit.
311set(_ADDL_TEST_SOURCES)
312if(MLIR_INCLUDE_TESTS)
313  set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
314  declare_mlir_python_sources(MLIRPythonTestSources)
315  declare_mlir_python_sources(MLIRPythonTestSources.Dialects
316    ADD_TO_PARENT MLIRPythonTestSources)
317
318  # TODO: this uses a tablegen file from the test directory and should be
319  # decoupled from here.
320  declare_mlir_python_sources(
321    MLIRPythonTestSources.Dialects.PythonTest
322    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
323    ADD_TO_PARENT MLIRPythonTestSources.Dialects
324    SOURCES dialects/python_test.py)
325  set(LLVM_TARGET_DEFINITIONS
326    "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
327  mlir_tablegen(
328    "dialects/_python_test_ops_gen.py"
329    -gen-python-op-bindings
330    -bind-dialect=python_test)
331  add_public_tablegen_target(PythonTestDialectPyIncGen)
332  declare_mlir_python_sources(
333    MLIRPythonTestSources.Dialects.PythonTest.ops_gen
334    ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
335    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
336    SOURCES "dialects/_python_test_ops_gen.py")
337
338  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
339    MODULE_NAME _mlirPythonTest
340    ADD_TO_PARENT MLIRPythonTestSources.Dialects
341    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
342    SOURCES
343      PythonTestModule.cpp
344    PRIVATE_LINK_LIBS
345      LLVMSupport
346    EMBED_CAPI_LINK_LIBS
347      MLIRCAPIPythonTestDialect
348  )
349endif()
350
351################################################################################
352# Common CAPI dependency DSO.
353# All python extensions must link through one DSO which exports the CAPI, and
354# this must have a globally unique name amongst all embeddors of the python
355# library since it will effectively have global scope.
356#
357# The presence of this aggregate library is part of the long term plan, but its
358# use needs to be made more flexible.
359#
360# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
361# once ready.
362################################################################################
363
364add_mlir_python_common_capi_library(MLIRPythonCAPI
365  INSTALL_COMPONENT MLIRPythonModules
366  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
367  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
368  RELATIVE_INSTALL_ROOT "../../../.."
369  DECLARED_SOURCES
370    MLIRPythonSources
371    MLIRPythonExtension.AllPassesRegistration
372    ${_ADDL_TEST_SOURCES}
373)
374
375################################################################################
376# The fully assembled package of modules.
377# This must come last.
378################################################################################
379
380add_mlir_python_modules(MLIRPythonModules
381  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
382  INSTALL_PREFIX "python_packages/mlir_core/mlir"
383  DECLARED_SOURCES
384    MLIRPythonSources
385    MLIRPythonExtension.AllPassesRegistration
386    MLIRPythonCAPIHeaderSources
387    ${_ADDL_TEST_SOURCES}
388  COMMON_CAPI_LINK_LIBS
389    MLIRPythonCAPI
390  )
391