xref: /llvm-project/mlir/python/CMakeLists.txt (revision 5cd427477218d8bdb659c6c53a7758f741c3990a)
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)
10declare_mlir_python_sources(MLIRPythonSources.Core
11  ADD_TO_PARENT MLIRPythonSources)
12
13################################################################################
14# Pure python sources and generated code
15################################################################################
16
17declare_mlir_python_sources(MLIRPythonSources.Core.Python
18  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
19  ADD_TO_PARENT MLIRPythonSources.Core
20  SOURCES
21    _mlir_libs/__init__.py
22    ir.py
23    passmanager.py
24    rewrite.py
25    dialects/_ods_common.py
26
27    # The main _mlir module has submodules: include stubs from each.
28    _mlir_libs/_mlir/__init__.pyi
29    _mlir_libs/_mlir/ir.pyi
30    _mlir_libs/_mlir/passmanager.pyi
31)
32
33declare_mlir_python_sources(MLIRPythonSources.Core.Python.Extras
34  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
35  ADD_TO_PARENT MLIRPythonSources.Core.Python
36  SOURCES
37    extras/types.py
38    extras/meta.py
39)
40
41declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
42  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
43  ADD_TO_PARENT MLIRPythonSources
44  SOURCES
45    execution_engine.py
46    _mlir_libs/_mlirExecutionEngine.pyi
47  SOURCES_GLOB
48    runtime/*.py
49)
50
51declare_mlir_python_sources(MLIRPythonCAPI.HeaderSources
52  ROOT_DIR "${MLIR_SOURCE_DIR}/include"
53  SOURCES_GLOB "mlir-c/*.h"
54)
55
56################################################################################
57# Dialect bindings
58################################################################################
59
60declare_mlir_dialect_python_bindings(
61  ADD_TO_PARENT MLIRPythonSources.Dialects
62  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
63  TD_FILE dialects/AffineOps.td
64  SOURCES
65    dialects/affine.py
66  DIALECT_NAME affine
67  GEN_ENUM_BINDINGS)
68
69declare_mlir_dialect_python_bindings(
70  ADD_TO_PARENT MLIRPythonSources.Dialects
71  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
72  TD_FILE dialects/AMDGPUOps.td
73  SOURCES
74    dialects/amdgpu.py
75  DIALECT_NAME amdgpu
76  GEN_ENUM_BINDINGS)
77
78declare_mlir_dialect_python_bindings(
79  ADD_TO_PARENT MLIRPythonSources.Dialects
80  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
81  TD_FILE dialects/AsyncOps.td
82  SOURCES_GLOB dialects/async_dialect/*.py
83  DIALECT_NAME async)
84
85declare_mlir_dialect_python_bindings(
86  ADD_TO_PARENT MLIRPythonSources.Dialects
87  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
88  TD_FILE dialects/BufferizationOps.td
89  SOURCES
90    dialects/bufferization.py
91  DIALECT_NAME bufferization
92  GEN_ENUM_BINDINGS_TD_FILE
93    "../../include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td"
94)
95
96declare_mlir_dialect_python_bindings(
97  ADD_TO_PARENT MLIRPythonSources.Dialects
98  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
99  TD_FILE dialects/BuiltinOps.td
100  SOURCES
101    dialects/builtin.py
102  DIALECT_NAME builtin)
103
104declare_mlir_dialect_python_bindings(
105  ADD_TO_PARENT MLIRPythonSources.Dialects
106  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
107  TD_FILE dialects/ComplexOps.td
108  SOURCES
109    dialects/complex.py
110  DIALECT_NAME complex)
111
112declare_mlir_dialect_python_bindings(
113  ADD_TO_PARENT MLIRPythonSources.Dialects
114  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
115  TD_FILE dialects/IndexOps.td
116  SOURCES
117    dialects/index.py
118  DIALECT_NAME index
119  GEN_ENUM_BINDINGS)
120
121declare_mlir_dialect_python_bindings(
122  ADD_TO_PARENT MLIRPythonSources.Dialects
123  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
124  TD_FILE dialects/ControlFlowOps.td
125  SOURCES
126    dialects/cf.py
127  DIALECT_NAME cf)
128
129declare_mlir_dialect_python_bindings(
130  ADD_TO_PARENT MLIRPythonSources.Dialects
131  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
132  TD_FILE dialects/FuncOps.td
133  SOURCES
134    dialects/func.py
135  DIALECT_NAME func)
136
137declare_mlir_dialect_python_bindings(
138  ADD_TO_PARENT MLIRPythonSources.Dialects
139  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
140  TD_FILE dialects/GPUOps.td
141  SOURCES_GLOB dialects/gpu/*.py
142  DIALECT_NAME gpu
143  GEN_ENUM_BINDINGS)
144
145declare_mlir_dialect_python_bindings(
146  ADD_TO_PARENT MLIRPythonSources.Dialects
147  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
148  TD_FILE dialects/LinalgOps.td
149  SOURCES
150  SOURCES_GLOB
151    dialects/linalg/*.py
152  DIALECT_NAME linalg
153  DEPENDS LinalgOdsGen
154  GEN_ENUM_BINDINGS)
155
156declare_mlir_dialect_python_bindings(
157  ADD_TO_PARENT MLIRPythonSources.Dialects
158  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
159  TD_FILE dialects/LLVMOps.td
160  SOURCES
161    dialects/llvm.py
162  DIALECT_NAME llvm
163  GEN_ENUM_BINDINGS)
164
165declare_mlir_dialect_extension_python_bindings(
166ADD_TO_PARENT MLIRPythonSources.Dialects
167ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
168  TD_FILE dialects/TransformPDLExtensionOps.td
169  SOURCES
170    dialects/transform/pdl.py
171  DIALECT_NAME transform
172  EXTENSION_NAME transform_pdl_extension)
173
174declare_mlir_dialect_python_bindings(
175  ADD_TO_PARENT MLIRPythonSources.Dialects
176  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
177  TD_FILE dialects/TransformOps.td
178  SOURCES
179    dialects/transform/__init__.py
180    _mlir_libs/_mlir/dialects/transform/__init__.pyi
181  DIALECT_NAME transform
182  GEN_ENUM_BINDINGS_TD_FILE
183    "../../include/mlir/Dialect/Transform/IR/TransformAttrs.td"
184)
185
186declare_mlir_python_sources(
187  MLIRPythonSources.Dialects.transform.extras
188  ADD_TO_PARENT MLIRPythonSources.Dialects
189  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
190  GEN_ENUM_BINDINGS
191  SOURCES
192    dialects/transform/extras/__init__.py)
193
194declare_mlir_python_sources(
195  MLIRPythonSources.Dialects.transform.interpreter
196  ADD_TO_PARENT MLIRPythonSources.Dialects
197  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
198  SOURCES
199    dialects/transform/interpreter/__init__.py)
200
201declare_mlir_dialect_extension_python_bindings(
202  ADD_TO_PARENT MLIRPythonSources.Dialects
203  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
204  TD_FILE dialects/BufferizationTransformOps.td
205  SOURCES
206    dialects/transform/bufferization.py
207  DIALECT_NAME transform
208  EXTENSION_NAME bufferization_transform)
209
210declare_mlir_dialect_extension_python_bindings(
211  ADD_TO_PARENT MLIRPythonSources.Dialects
212  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
213  TD_FILE dialects/GPUTransformOps.td
214  SOURCES
215    dialects/transform/gpu.py
216  DIALECT_NAME transform
217  EXTENSION_NAME gpu_transform)
218
219declare_mlir_dialect_extension_python_bindings(
220  ADD_TO_PARENT MLIRPythonSources.Dialects
221  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
222  TD_FILE dialects/SCFLoopTransformOps.td
223  SOURCES
224    dialects/transform/loop.py
225  DIALECT_NAME transform
226  EXTENSION_NAME loop_transform)
227
228declare_mlir_dialect_extension_python_bindings(
229  ADD_TO_PARENT MLIRPythonSources.Dialects
230  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
231  TD_FILE dialects/MemRefTransformOps.td
232  SOURCES
233    dialects/transform/memref.py
234  DIALECT_NAME transform
235  EXTENSION_NAME memref_transform)
236
237declare_mlir_dialect_extension_python_bindings(
238  ADD_TO_PARENT MLIRPythonSources.Dialects
239  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
240  TD_FILE dialects/NVGPUTransformOps.td
241  SOURCES
242    dialects/transform/nvgpu.py
243  DIALECT_NAME transform
244  EXTENSION_NAME nvgpu_transform)
245
246declare_mlir_dialect_extension_python_bindings(
247  ADD_TO_PARENT MLIRPythonSources.Dialects
248  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
249  TD_FILE dialects/LinalgStructuredTransformOps.td
250  SOURCES
251    dialects/transform/structured.py
252  DIALECT_NAME transform
253  EXTENSION_NAME structured_transform
254  GEN_ENUM_BINDINGS_TD_FILE
255    "../../include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td"
256)
257
258declare_mlir_dialect_extension_python_bindings(
259  ADD_TO_PARENT MLIRPythonSources.Dialects
260  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
261  TD_FILE dialects/SparseTensorTransformOps.td
262  SOURCES
263    dialects/transform/sparse_tensor.py
264  DIALECT_NAME transform
265  EXTENSION_NAME sparse_tensor_transform)
266
267declare_mlir_dialect_extension_python_bindings(
268  ADD_TO_PARENT MLIRPythonSources.Dialects
269  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
270  TD_FILE dialects/TensorTransformOps.td
271  SOURCES
272    dialects/transform/tensor.py
273  DIALECT_NAME transform
274  EXTENSION_NAME tensor_transform)
275
276declare_mlir_dialect_extension_python_bindings(
277  ADD_TO_PARENT MLIRPythonSources.Dialects
278  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
279  TD_FILE dialects/VectorTransformOps.td
280  SOURCES
281    dialects/transform/vector.py
282  DIALECT_NAME transform
283  EXTENSION_NAME vector_transform
284  GEN_ENUM_BINDINGS_TD_FILE
285    "../../include/mlir/Dialect/Vector/Transforms/VectorTransformsBase.td"
286)
287
288declare_mlir_dialect_python_bindings(
289  ADD_TO_PARENT MLIRPythonSources.Dialects
290  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
291  TD_FILE dialects/MathOps.td
292  SOURCES dialects/math.py
293  DIALECT_NAME math)
294
295declare_mlir_dialect_python_bindings(
296  ADD_TO_PARENT MLIRPythonSources.Dialects
297  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
298  TD_FILE dialects/ArithOps.td
299  SOURCES
300    dialects/arith.py
301  DIALECT_NAME arith
302  GEN_ENUM_BINDINGS)
303
304declare_mlir_dialect_python_bindings(
305  ADD_TO_PARENT MLIRPythonSources.Dialects
306  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
307  TD_FILE dialects/MemRefOps.td
308  SOURCES
309    dialects/memref.py
310  DIALECT_NAME memref)
311
312declare_mlir_dialect_python_bindings(
313  ADD_TO_PARENT MLIRPythonSources.Dialects
314  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
315  TD_FILE dialects/MLProgramOps.td
316  SOURCES
317    dialects/ml_program.py
318  DIALECT_NAME ml_program)
319
320declare_mlir_dialect_python_bindings(
321  ADD_TO_PARENT MLIRPythonSources.Dialects
322  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
323  TD_FILE dialects/NVGPUOps.td
324  SOURCES
325    dialects/nvgpu.py
326  DIALECT_NAME nvgpu
327  GEN_ENUM_BINDINGS)
328
329declare_mlir_dialect_python_bindings(
330  ADD_TO_PARENT MLIRPythonSources.Dialects
331  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
332  TD_FILE dialects/NVVMOps.td
333  SOURCES
334    dialects/nvvm.py
335  DIALECT_NAME nvvm
336  GEN_ENUM_BINDINGS)
337
338declare_mlir_dialect_python_bindings(
339  ADD_TO_PARENT MLIRPythonSources.Dialects
340  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
341  TD_FILE dialects/ROCDLOps.td
342  SOURCES
343    dialects/rocdl.py
344  DIALECT_NAME rocdl)
345
346declare_mlir_python_sources(
347  MLIRPythonSources.Dialects.quant
348  ADD_TO_PARENT MLIRPythonSources.Dialects
349  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
350  GEN_ENUM_BINDINGS
351  SOURCES
352    dialects/quant.py
353    _mlir_libs/_mlir/dialects/quant.pyi)
354
355declare_mlir_dialect_python_bindings(
356  ADD_TO_PARENT MLIRPythonSources.Dialects
357  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
358  TD_FILE dialects/EmitC.td
359  SOURCES
360    dialects/emitc.py
361  DIALECT_NAME emitc)
362
363declare_mlir_dialect_python_bindings(
364  ADD_TO_PARENT MLIRPythonSources.Dialects
365  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
366  TD_FILE dialects/PDLOps.td
367  SOURCES
368    dialects/pdl.py
369    _mlir_libs/_mlir/dialects/pdl.pyi
370  DIALECT_NAME pdl)
371
372declare_mlir_dialect_python_bindings(
373  ADD_TO_PARENT MLIRPythonSources.Dialects
374  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
375  TD_FILE dialects/OpenMPOps.td
376  SOURCES
377    dialects/openmp.py
378  DIALECT_NAME omp
379  DEPENDS omp_common_td)
380
381declare_mlir_dialect_python_bindings(
382  ADD_TO_PARENT MLIRPythonSources.Dialects
383  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
384  TD_FILE dialects/SCFOps.td
385  SOURCES
386    dialects/scf.py
387  DIALECT_NAME scf)
388
389declare_mlir_dialect_python_bindings(
390  ADD_TO_PARENT MLIRPythonSources.Dialects
391  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
392  TD_FILE dialects/ShapeOps.td
393  SOURCES dialects/shape.py
394  DIALECT_NAME shape)
395
396declare_mlir_dialect_python_bindings(
397  ADD_TO_PARENT MLIRPythonSources.Dialects
398  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
399  TD_FILE dialects/SparseTensorOps.td
400  SOURCES dialects/sparse_tensor.py
401  DIALECT_NAME sparse_tensor
402  GEN_ENUM_BINDINGS_TD_FILE
403    "../../include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td"
404)
405
406declare_mlir_dialect_python_bindings(
407    ADD_TO_PARENT MLIRPythonSources.Dialects
408    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
409    TD_FILE dialects/SPIRVOps.td
410    SOURCES dialects/spirv.py
411    DIALECT_NAME spirv)
412
413declare_mlir_dialect_python_bindings(
414  ADD_TO_PARENT MLIRPythonSources.Dialects
415  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
416  TD_FILE dialects/TensorOps.td
417  SOURCES
418    dialects/tensor.py
419  DIALECT_NAME tensor)
420
421declare_mlir_dialect_python_bindings(
422  ADD_TO_PARENT MLIRPythonSources.Dialects
423  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
424  TD_FILE dialects/TosaOps.td
425  SOURCES dialects/tosa.py
426  DIALECT_NAME tosa
427)
428
429declare_mlir_dialect_python_bindings(
430  ADD_TO_PARENT MLIRPythonSources.Dialects
431  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
432  TD_FILE dialects/VectorOps.td
433  SOURCES dialects/vector.py
434  DIALECT_NAME vector
435  GEN_ENUM_BINDINGS_TD_FILE
436    "dialects/VectorAttributes.td")
437
438################################################################################
439# Python extensions.
440# The sources for these are all in lib/Bindings/Python, but since they have to
441# be rebuilt for each package and integrate with the source setup here, we
442# just reference them here instead of having ordered, cross package target
443# dependencies.
444################################################################################
445
446set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
447declare_mlir_python_extension(MLIRPythonExtension.Core
448  MODULE_NAME _mlir
449  ADD_TO_PARENT MLIRPythonSources.Core
450  ROOT_DIR "${PYTHON_SOURCE_DIR}"
451  PYTHON_BINDINGS_LIBRARY nanobind
452  SOURCES
453    MainModule.cpp
454    IRAffine.cpp
455    IRAttributes.cpp
456    IRCore.cpp
457    IRInterfaces.cpp
458    IRModule.cpp
459    IRTypes.cpp
460    Pass.cpp
461    Rewrite.cpp
462
463    # Headers must be included explicitly so they are installed.
464    Globals.h
465    IRModule.h
466    Pass.h
467    NanobindUtils.h
468    Rewrite.h
469  PRIVATE_LINK_LIBS
470    LLVMSupport
471  EMBED_CAPI_LINK_LIBS
472    MLIRCAPIDebug
473    MLIRCAPIIR
474    MLIRCAPIInterfaces
475
476    # Dialects
477    MLIRCAPIFunc
478)
479
480# This extension exposes an API to register all dialects, extensions, and passes
481# packaged in upstream MLIR and it is used for the upstream "mlir" Python
482# package. Downstreams will likely want to provide their own and not depend
483# on this one, since it links in the world.
484# Note that this is not added to any top-level source target for transitive
485# inclusion: It must be included explicitly by downstreams if desired. Note that
486# this has a very large impact on what gets built/packaged.
487declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
488  MODULE_NAME _mlirRegisterEverything
489  ROOT_DIR "${PYTHON_SOURCE_DIR}"
490  PYTHON_BINDINGS_LIBRARY nanobind
491  SOURCES
492    RegisterEverything.cpp
493  PRIVATE_LINK_LIBS
494    LLVMSupport
495  EMBED_CAPI_LINK_LIBS
496    MLIRCAPIConversion
497    MLIRCAPITransforms
498    MLIRCAPIRegisterEverything
499)
500
501declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
502  MODULE_NAME _mlirDialectsLinalg
503  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
504  ROOT_DIR "${PYTHON_SOURCE_DIR}"
505  PYTHON_BINDINGS_LIBRARY nanobind
506  SOURCES
507    DialectLinalg.cpp
508  PRIVATE_LINK_LIBS
509    LLVMSupport
510  EMBED_CAPI_LINK_LIBS
511    MLIRCAPIIR
512    MLIRCAPILinalg
513)
514
515declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
516  MODULE_NAME _mlirDialectsGPU
517  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
518  ROOT_DIR "${PYTHON_SOURCE_DIR}"
519  PYTHON_BINDINGS_LIBRARY nanobind
520  SOURCES
521    DialectGPU.cpp
522  PRIVATE_LINK_LIBS
523    LLVMSupport
524  EMBED_CAPI_LINK_LIBS
525    MLIRCAPIIR
526    MLIRCAPIGPU
527)
528
529declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
530  MODULE_NAME _mlirDialectsLLVM
531  ADD_TO_PARENT MLIRPythonSources.Dialects.llvm
532  ROOT_DIR "${PYTHON_SOURCE_DIR}"
533  PYTHON_BINDINGS_LIBRARY nanobind
534  SOURCES
535    DialectLLVM.cpp
536  PRIVATE_LINK_LIBS
537    LLVMSupport
538  EMBED_CAPI_LINK_LIBS
539    MLIRCAPIIR
540    MLIRCAPILLVM
541)
542
543declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
544  MODULE_NAME _mlirDialectsQuant
545  ADD_TO_PARENT MLIRPythonSources.Dialects.quant
546  ROOT_DIR "${PYTHON_SOURCE_DIR}"
547  PYTHON_BINDINGS_LIBRARY nanobind
548  SOURCES
549    DialectQuant.cpp
550  PRIVATE_LINK_LIBS
551    LLVMSupport
552  EMBED_CAPI_LINK_LIBS
553    MLIRCAPIIR
554    MLIRCAPIQuant
555)
556
557declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
558  MODULE_NAME _mlirDialectsNVGPU
559  ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu
560  ROOT_DIR "${PYTHON_SOURCE_DIR}"
561  PYTHON_BINDINGS_LIBRARY nanobind
562  SOURCES
563    DialectNVGPU.cpp
564  PRIVATE_LINK_LIBS
565    LLVMSupport
566  EMBED_CAPI_LINK_LIBS
567    MLIRCAPIIR
568    MLIRCAPINVGPU
569)
570
571declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
572  MODULE_NAME _mlirDialectsPDL
573  ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
574  ROOT_DIR "${PYTHON_SOURCE_DIR}"
575  PYTHON_BINDINGS_LIBRARY nanobind
576  SOURCES
577    DialectPDL.cpp
578  PRIVATE_LINK_LIBS
579    LLVMSupport
580  EMBED_CAPI_LINK_LIBS
581    MLIRCAPIIR
582    MLIRCAPIPDL
583)
584
585declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
586  MODULE_NAME _mlirDialectsSparseTensor
587  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
588  ROOT_DIR "${PYTHON_SOURCE_DIR}"
589  PYTHON_BINDINGS_LIBRARY nanobind
590  SOURCES
591    DialectSparseTensor.cpp
592  PRIVATE_LINK_LIBS
593    LLVMSupport
594  EMBED_CAPI_LINK_LIBS
595    MLIRCAPIIR
596    MLIRCAPISparseTensor
597)
598
599declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
600  MODULE_NAME _mlirDialectsTransform
601  ADD_TO_PARENT MLIRPythonSources.Dialects.transform
602  ROOT_DIR "${PYTHON_SOURCE_DIR}"
603  PYTHON_BINDINGS_LIBRARY nanobind
604  SOURCES
605    DialectTransform.cpp
606  PRIVATE_LINK_LIBS
607    LLVMSupport
608  EMBED_CAPI_LINK_LIBS
609    MLIRCAPIIR
610    MLIRCAPITransformDialect
611)
612
613declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
614  MODULE_NAME _mlirAsyncPasses
615  ADD_TO_PARENT MLIRPythonSources.Dialects.async
616  ROOT_DIR "${PYTHON_SOURCE_DIR}"
617  PYTHON_BINDINGS_LIBRARY nanobind
618  SOURCES
619    AsyncPasses.cpp
620  PRIVATE_LINK_LIBS
621    LLVMSupport
622  EMBED_CAPI_LINK_LIBS
623    MLIRCAPIAsync
624)
625
626if(MLIR_ENABLE_EXECUTION_ENGINE)
627  declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
628    MODULE_NAME _mlirExecutionEngine
629    ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
630    ROOT_DIR "${PYTHON_SOURCE_DIR}"
631  PYTHON_BINDINGS_LIBRARY nanobind
632    SOURCES
633      ExecutionEngineModule.cpp
634    PRIVATE_LINK_LIBS
635      LLVMSupport
636    EMBED_CAPI_LINK_LIBS
637      MLIRCAPIExecutionEngine
638  )
639endif()
640
641declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
642  MODULE_NAME _mlirGPUPasses
643  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
644  ROOT_DIR "${PYTHON_SOURCE_DIR}"
645  PYTHON_BINDINGS_LIBRARY nanobind
646  SOURCES
647    GPUPasses.cpp
648  PRIVATE_LINK_LIBS
649    LLVMSupport
650  EMBED_CAPI_LINK_LIBS
651    MLIRCAPIGPU
652)
653
654declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
655  MODULE_NAME _mlirLinalgPasses
656  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
657  ROOT_DIR "${PYTHON_SOURCE_DIR}"
658  PYTHON_BINDINGS_LIBRARY nanobind
659  SOURCES
660    LinalgPasses.cpp
661  PRIVATE_LINK_LIBS
662    LLVMSupport
663  EMBED_CAPI_LINK_LIBS
664    MLIRCAPILinalg
665)
666
667declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
668  MODULE_NAME _mlirSparseTensorPasses
669  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
670  ROOT_DIR "${PYTHON_SOURCE_DIR}"
671  PYTHON_BINDINGS_LIBRARY nanobind
672  SOURCES
673    SparseTensorPasses.cpp
674  PRIVATE_LINK_LIBS
675    LLVMSupport
676  EMBED_CAPI_LINK_LIBS
677    MLIRCAPISparseTensor
678)
679
680declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
681  MODULE_NAME _mlirTransformInterpreter
682  ADD_TO_PARENT MLIRPythonSources.Dialects.transform
683  ROOT_DIR "${PYTHON_SOURCE_DIR}"
684  PYTHON_BINDINGS_LIBRARY nanobind
685  SOURCES
686    TransformInterpreter.cpp
687  PRIVATE_LINK_LIBS
688    LLVMSupport
689  EMBED_CAPI_LINK_LIBS
690    MLIRCAPITransformDialectTransforms
691)
692
693# TODO: Figure out how to put this in the test tree.
694# This should not be included in the main Python extension. However,
695# putting it into MLIRPythonTestSources along with the dialect declaration
696# above confuses Python module loader when running under lit.
697set(_ADDL_TEST_SOURCES)
698if(MLIR_INCLUDE_TESTS)
699  set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
700  declare_mlir_python_sources(MLIRPythonTestSources)
701  declare_mlir_python_sources(MLIRPythonTestSources.Dialects
702    ADD_TO_PARENT MLIRPythonTestSources)
703
704  # TODO: this uses a tablegen file from the test directory and should be
705  # decoupled from here.
706  declare_mlir_python_sources(
707    MLIRPythonTestSources.Dialects.PythonTest
708    ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
709    ADD_TO_PARENT MLIRPythonTestSources.Dialects
710    SOURCES
711      dialects/python_test.py
712  )
713  set(LLVM_TARGET_DEFINITIONS
714    "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
715  mlir_tablegen(
716    "dialects/_python_test_ops_gen.py"
717    -gen-python-op-bindings
718    -bind-dialect=python_test)
719  add_public_tablegen_target(PythonTestDialectPyIncGen)
720  declare_mlir_python_sources(
721    MLIRPythonTestSources.Dialects.PythonTest.ops_gen
722    ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
723    ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
724    SOURCES "dialects/_python_test_ops_gen.py")
725
726  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionPybind11
727    MODULE_NAME _mlirPythonTestPybind11
728    ADD_TO_PARENT MLIRPythonTestSources.Dialects
729    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
730    PYTHON_BINDINGS_LIBRARY pybind11
731    SOURCES
732      PythonTestModulePybind11.cpp
733    PRIVATE_LINK_LIBS
734      LLVMSupport
735    EMBED_CAPI_LINK_LIBS
736      MLIRCAPIPythonTestDialect
737  )
738  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionNanobind
739    MODULE_NAME _mlirPythonTestNanobind
740    ADD_TO_PARENT MLIRPythonTestSources.Dialects
741    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
742    PYTHON_BINDINGS_LIBRARY nanobind
743    SOURCES
744      PythonTestModuleNanobind.cpp
745    PRIVATE_LINK_LIBS
746      LLVMSupport
747    EMBED_CAPI_LINK_LIBS
748      MLIRCAPIPythonTestDialect
749  )
750endif()
751
752################################################################################
753# Common CAPI dependency DSO.
754# All python extensions must link through one DSO which exports the CAPI, and
755# this must have a globally unique name amongst all embeddors of the python
756# library since it will effectively have global scope.
757#
758# The presence of this aggregate library is part of the long term plan, but its
759# use needs to be made more flexible.
760#
761# TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
762# once ready.
763################################################################################
764
765add_mlir_python_common_capi_library(MLIRPythonCAPI
766  INSTALL_COMPONENT MLIRPythonModules
767  INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
768  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
769  RELATIVE_INSTALL_ROOT "../../../.."
770  DECLARED_HEADERS
771    MLIRPythonCAPI.HeaderSources
772  DECLARED_SOURCES
773    MLIRPythonSources
774    MLIRPythonExtension.RegisterEverything
775    ${_ADDL_TEST_SOURCES}
776)
777
778################################################################################
779# Custom targets.
780################################################################################
781
782_flatten_mlir_python_targets(mlir_python_sources_deps MLIRPythonSources)
783add_custom_target("mlir-python-sources" DEPENDS ${mlir_python_sources_deps})
784if(NOT LLVM_ENABLE_IDE)
785  add_llvm_install_targets(install-mlir-python-sources
786    DEPENDS mlir-python-sources
787    COMPONENT mlir-python-sources
788  )
789endif()
790
791################################################################################
792# The fully assembled package of modules.
793# This must come last.
794################################################################################
795
796add_mlir_python_modules(MLIRPythonModules
797  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
798  INSTALL_PREFIX "python_packages/mlir_core/mlir"
799  DECLARED_SOURCES
800    MLIRPythonSources
801    MLIRPythonExtension.RegisterEverything
802    ${_ADDL_TEST_SOURCES}
803  COMMON_CAPI_LINK_LIBS
804    MLIRPythonCAPI
805)
806
807