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  SOURCES
178    ${PYTHON_SOURCE_DIR}/DialectLinalg.cpp  # TODO: Break this out.
179    ${PYTHON_SOURCE_DIR}/DialectSparseTensor.cpp  # TODO: Break this out.
180    ${PYTHON_SOURCE_DIR}/MainModule.cpp
181    ${PYTHON_SOURCE_DIR}/IRAffine.cpp
182    ${PYTHON_SOURCE_DIR}/IRAttributes.cpp
183    ${PYTHON_SOURCE_DIR}/IRCore.cpp
184    ${PYTHON_SOURCE_DIR}/IRInterfaces.cpp
185    ${PYTHON_SOURCE_DIR}/IRModule.cpp
186    ${PYTHON_SOURCE_DIR}/IRTypes.cpp
187    ${PYTHON_SOURCE_DIR}/PybindUtils.cpp
188    ${PYTHON_SOURCE_DIR}/Pass.cpp
189  PRIVATE_LINK_LIBS
190    LLVMSupport
191  EMBED_CAPI_LINK_LIBS
192    MLIRCAPIDebug
193    MLIRCAPIIR
194    MLIRCAPIInterfaces
195    MLIRCAPIRegistration  # TODO: See about dis-aggregating
196
197    # Dialects
198    MLIRCAPILinalg  # TODO: Remove when above is removed.
199    MLIRCAPISparseTensor  # TODO: Remove when above is removed.
200    MLIRCAPIStandard
201)
202
203declare_mlir_python_extension(MLIRPythonExtension.AllPassesRegistration
204  MODULE_NAME _mlirAllPassesRegistration
205  SOURCES
206    ${PYTHON_SOURCE_DIR}/AllPassesRegistration.cpp
207  PRIVATE_LINK_LIBS
208    LLVMSupport
209  EMBED_CAPI_LINK_LIBS
210    MLIRCAPIConversion
211    MLIRCAPITransforms
212)
213
214declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
215  MODULE_NAME _mlirAsyncPasses
216  ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
217  SOURCES
218    ${PYTHON_SOURCE_DIR}/AsyncPasses.cpp
219  PRIVATE_LINK_LIBS
220    LLVMSupport
221  EMBED_CAPI_LINK_LIBS
222    MLIRCAPIAsync
223)
224
225declare_mlir_python_extension(MLIRPythonExtension.Conversions
226  MODULE_NAME _mlirConversions
227  ADD_TO_PARENT MLIRPythonSources.Passes
228  SOURCES
229    ${PYTHON_SOURCE_DIR}/Conversions/Conversions.cpp
230  PRIVATE_LINK_LIBS
231    LLVMSupport
232  EMBED_CAPI_LINK_LIBS
233  MLIRCAPIConversion
234)
235
236declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
237  MODULE_NAME _mlirExecutionEngine
238  ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
239  SOURCES
240    ${PYTHON_SOURCE_DIR}/ExecutionEngineModule.cpp
241  PRIVATE_LINK_LIBS
242    LLVMSupport
243  EMBED_CAPI_LINK_LIBS
244    MLIRCEXECUTIONENGINE
245)
246
247declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
248  MODULE_NAME _mlirGPUPasses
249  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
250  SOURCES
251    ${PYTHON_SOURCE_DIR}/GPUPasses.cpp
252  PRIVATE_LINK_LIBS
253    LLVMSupport
254  EMBED_CAPI_LINK_LIBS
255    MLIRCAPIGPU
256)
257
258declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
259  MODULE_NAME _mlirLinalgPasses
260  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
261  SOURCES
262    ${PYTHON_SOURCE_DIR}/LinalgPasses.cpp
263  PRIVATE_LINK_LIBS
264    LLVMSupport
265  EMBED_CAPI_LINK_LIBS
266    MLIRCAPILinalg
267)
268
269declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
270  MODULE_NAME _mlirSparseTensorPasses
271  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
272  SOURCES
273    ${PYTHON_SOURCE_DIR}/SparseTensorPasses.cpp
274  PRIVATE_LINK_LIBS
275    LLVMSupport
276  EMBED_CAPI_LINK_LIBS
277    MLIRCAPISparseTensor
278)
279
280declare_mlir_python_extension(MLIRPythonExtension.Transforms
281  MODULE_NAME _mlirTransforms
282  ADD_TO_PARENT MLIRPythonSources.Passes
283  SOURCES
284    ${PYTHON_SOURCE_DIR}/Transforms/Transforms.cpp
285  PRIVATE_LINK_LIBS
286    LLVMSupport
287  EMBED_CAPI_LINK_LIBS
288    MLIRCAPITransforms
289)
290
291# TODO: Figure out how to put this in the test tree.
292# This should not be included in the main Python extension. However,
293# putting it into MLIRPythonTestSources along with the dialect declaration
294# above confuses Python module loader when running under lit.
295set(_ADDL_TEST_SOURCES)
296if(MLIR_INCLUDE_TESTS)
297  set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
298  declare_mlir_python_sources(MLIRPythonTestSources)
299  declare_mlir_python_sources(MLIRPythonTestSources.Dialects
300    ADD_TO_PARENT MLIRPythonTestSources)
301
302  # TODO: this uses a tablegen file from the test directory and should be
303  # decoupled from here.
304  declare_mlir_python_sources(
305    MLIRPythonTestSources.Dialects.PythonTest
306    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
307    ADD_TO_PARENT MLIRPythonTestSources.Dialects
308    SOURCES dialects/python_test.py)
309  set(LLVM_TARGET_DEFINITIONS
310    "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
311  mlir_tablegen(
312    "dialects/_python_test_ops_gen.py"
313    -gen-python-op-bindings
314    -bind-dialect=python_test)
315  add_public_tablegen_target(PythonTestDialectPyIncGen)
316  declare_mlir_python_sources(
317    MLIRPythonTestSources.Dialects.PythonTest.ops_gen
318    ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
319    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
320    SOURCES "dialects/_python_test_ops_gen.py")
321
322  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
323    MODULE_NAME _mlirPythonTest
324    ADD_TO_PARENT MLIRPythonTestSources.Dialects
325    SOURCES
326      ${MLIR_SOURCE_DIR}/test/python/lib/PythonTestModule.cpp
327    PRIVATE_LINK_LIBS
328      LLVMSupport
329    EMBED_CAPI_LINK_LIBS
330      MLIRCAPIPythonTestDialect
331  )
332endif()
333
334################################################################################
335# Common CAPI dependency DSO.
336# All python extensions must link through one DSO which exports the CAPI, and
337# this must have a globally unique name amongst all embeddors of the python
338# library since it will effectively have global scope.
339#
340# The presence of this aggregate library is part of the long term plan, but its
341# use needs to be made more flexible.
342#
343# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
344# once ready.
345################################################################################
346
347add_mlir_python_common_capi_library(MLIRPythonCAPI
348  INSTALL_COMPONENT MLIRPythonModules
349  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
350  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
351  RELATIVE_INSTALL_ROOT "../../../.."
352  DECLARED_SOURCES
353    MLIRPythonSources
354    MLIRPythonExtension.AllPassesRegistration
355    ${_ADDL_TEST_SOURCES}
356)
357
358################################################################################
359# The fully assembled package of modules.
360# This must come last.
361################################################################################
362
363add_mlir_python_modules(MLIRPythonModules
364  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
365  INSTALL_PREFIX "python_packages/mlir_core/mlir"
366  DECLARED_SOURCES
367    MLIRPythonSources
368    MLIRPythonExtension.AllPassesRegistration
369    MLIRPythonCAPIHeaderSources
370    ${_ADDL_TEST_SOURCES}
371  COMMON_CAPI_LINK_LIBS
372    MLIRPythonCAPI
373  )
374