mlir/python/CMakeLists.txt

   1 include(AddMLIRPython)
   2
   3 ################################################################################
   4 # Structural groupings.
   5 ################################################################################
   6
   7 declare_mlir_python_sources(MLIRPythonSources)
   8 declare_mlir_python_sources(MLIRPythonSources.Dialects
   9   ADD_TO_PARENT MLIRPythonSources)
  10 declare_mlir_python_sources(MLIRPythonSources.Core
  11   ADD_TO_PARENT MLIRPythonSources)
  12
  13 ################################################################################
  14 # Pure python sources and generated code
  15 ################################################################################
  16
  17 declare_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
  33 declare_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
  41 declare_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
  51 declare_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
  60 declare_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
  69 declare_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
  78 declare_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
  85 declare_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
  96 declare_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
 104 declare_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
 112 declare_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
 121 declare_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
 129 declare_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
 137 declare_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
 145 declare_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
 156 declare_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
 165 declare_mlir_dialect_extension_python_bindings(
 166 ADD_TO_PARENT MLIRPythonSources.Dialects
 167 ROOT_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
 174 declare_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
 186 declare_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
 194 declare_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
 201 declare_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
 210 declare_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
 219 declare_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
 228 declare_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
 237 declare_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
 246 declare_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
 258 declare_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
 267 declare_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
 276 declare_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
 288 declare_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
 295 declare_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
 304 declare_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
 312 declare_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
 320 declare_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
 329 declare_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
 338 declare_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
 346 declare_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
 355 declare_mlir_dialect_python_bindings(
 356   ADD_TO_PARENT MLIRPythonSources.Dialects
 357   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 358   TD_FILE dialects/PDLOps.td
 359   SOURCES
 360     dialects/pdl.py
 361     _mlir_libs/_mlir/dialects/pdl.pyi
 362   DIALECT_NAME pdl)
 363
 364 declare_mlir_dialect_python_bindings(
 365   ADD_TO_PARENT MLIRPythonSources.Dialects
 366   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 367   TD_FILE dialects/OpenMPOps.td
 368   SOURCES
 369     dialects/openmp.py
 370   DIALECT_NAME omp
 371   DEPENDS omp_common_td)
 372
 373 declare_mlir_dialect_python_bindings(
 374   ADD_TO_PARENT MLIRPythonSources.Dialects
 375   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 376   TD_FILE dialects/SCFOps.td
 377   SOURCES
 378     dialects/scf.py
 379   DIALECT_NAME scf)
 380
 381 declare_mlir_dialect_python_bindings(
 382   ADD_TO_PARENT MLIRPythonSources.Dialects
 383   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 384   TD_FILE dialects/ShapeOps.td
 385   SOURCES dialects/shape.py
 386   DIALECT_NAME shape)
 387
 388 declare_mlir_dialect_python_bindings(
 389   ADD_TO_PARENT MLIRPythonSources.Dialects
 390   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 391   TD_FILE dialects/SparseTensorOps.td
 392   SOURCES dialects/sparse_tensor.py
 393   DIALECT_NAME sparse_tensor
 394   GEN_ENUM_BINDINGS_TD_FILE
 395     "../../include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td"
 396 )
 397
 398 declare_mlir_dialect_python_bindings(
 399     ADD_TO_PARENT MLIRPythonSources.Dialects
 400     ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 401     TD_FILE dialects/SPIRVOps.td
 402     SOURCES dialects/spirv.py
 403     DIALECT_NAME spirv)
 404
 405 declare_mlir_dialect_python_bindings(
 406   ADD_TO_PARENT MLIRPythonSources.Dialects
 407   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 408   TD_FILE dialects/TensorOps.td
 409   SOURCES
 410     dialects/tensor.py
 411   DIALECT_NAME tensor)
 412
 413 declare_mlir_dialect_python_bindings(
 414   ADD_TO_PARENT MLIRPythonSources.Dialects
 415   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 416   TD_FILE dialects/TosaOps.td
 417   SOURCES dialects/tosa.py
 418   DIALECT_NAME tosa
 419 )
 420
 421 declare_mlir_dialect_python_bindings(
 422   ADD_TO_PARENT MLIRPythonSources.Dialects
 423   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 424   TD_FILE dialects/VectorOps.td
 425   SOURCES dialects/vector.py
 426   DIALECT_NAME vector
 427   GEN_ENUM_BINDINGS_TD_FILE
 428     "dialects/VectorAttributes.td")
 429
 430 ################################################################################
 431 # Python extensions.
 432 # The sources for these are all in lib/Bindings/Python, but since they have to
 433 # be rebuilt for each package and integrate with the source setup here, we
 434 # just reference them here instead of having ordered, cross package target
 435 # dependencies.
 436 ################################################################################
 437
 438 set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
 439 declare_mlir_python_extension(MLIRPythonExtension.Core
 440   MODULE_NAME _mlir
 441   ADD_TO_PARENT MLIRPythonSources.Core
 442   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 443   SOURCES
 444     MainModule.cpp
 445     IRAffine.cpp
 446     IRAttributes.cpp
 447     IRCore.cpp
 448     IRInterfaces.cpp
 449     IRModule.cpp
 450     IRTypes.cpp
 451     Pass.cpp
 452     Rewrite.cpp
 453
 454     # Headers must be included explicitly so they are installed.
 455     Globals.h
 456     IRModule.h
 457     Pass.h
 458     PybindUtils.h
 459     Rewrite.h
 460   PRIVATE_LINK_LIBS
 461     LLVMSupport
 462   EMBED_CAPI_LINK_LIBS
 463     MLIRCAPIDebug
 464     MLIRCAPIIR
 465     MLIRCAPIInterfaces
 466
 467     # Dialects
 468     MLIRCAPIFunc
 469 )
 470
 471 # This extension exposes an API to register all dialects, extensions, and passes
 472 # packaged in upstream MLIR and it is used for the upstream "mlir" Python
 473 # package. Downstreams will likely want to provide their own and not depend
 474 # on this one, since it links in the world.
 475 # Note that this is not added to any top-level source target for transitive
 476 # inclusion: It must be included explicitly by downstreams if desired. Note that
 477 # this has a very large impact on what gets built/packaged.
 478 declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
 479   MODULE_NAME _mlirRegisterEverything
 480   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 481   SOURCES
 482     RegisterEverything.cpp
 483   PRIVATE_LINK_LIBS
 484     LLVMSupport
 485   EMBED_CAPI_LINK_LIBS
 486     MLIRCAPIConversion
 487     MLIRCAPITransforms
 488     MLIRCAPIRegisterEverything
 489 )
 490
 491 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
 492   MODULE_NAME _mlirDialectsLinalg
 493   ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
 494   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 495   SOURCES
 496     DialectLinalg.cpp
 497   PRIVATE_LINK_LIBS
 498     LLVMSupport
 499   EMBED_CAPI_LINK_LIBS
 500     MLIRCAPIIR
 501     MLIRCAPILinalg
 502 )
 503
 504 declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
 505   MODULE_NAME _mlirDialectsGPU
 506   ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
 507   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 508   SOURCES
 509     DialectGPU.cpp
 510   PRIVATE_LINK_LIBS
 511     LLVMSupport
 512   EMBED_CAPI_LINK_LIBS
 513     MLIRCAPIIR
 514     MLIRCAPIGPU
 515 )
 516
 517 declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
 518   MODULE_NAME _mlirDialectsLLVM
 519   ADD_TO_PARENT MLIRPythonSources.Dialects.llvm
 520   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 521   SOURCES
 522     DialectLLVM.cpp
 523   PRIVATE_LINK_LIBS
 524     LLVMSupport
 525   EMBED_CAPI_LINK_LIBS
 526     MLIRCAPIIR
 527     MLIRCAPILLVM
 528 )
 529
 530 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
 531   MODULE_NAME _mlirDialectsQuant
 532   ADD_TO_PARENT MLIRPythonSources.Dialects.quant
 533   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 534   SOURCES
 535     DialectQuant.cpp
 536   PRIVATE_LINK_LIBS
 537     LLVMSupport
 538   EMBED_CAPI_LINK_LIBS
 539     MLIRCAPIIR
 540     MLIRCAPIQuant
 541 )
 542
 543 declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
 544   MODULE_NAME _mlirDialectsNVGPU
 545   ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu
 546   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 547   SOURCES
 548     DialectNVGPU.cpp
 549   PRIVATE_LINK_LIBS
 550     LLVMSupport
 551   EMBED_CAPI_LINK_LIBS
 552     MLIRCAPIIR
 553     MLIRCAPINVGPU
 554 )
 555
 556 declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
 557   MODULE_NAME _mlirDialectsPDL
 558   ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
 559   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 560   SOURCES
 561     DialectPDL.cpp
 562   PRIVATE_LINK_LIBS
 563     LLVMSupport
 564   EMBED_CAPI_LINK_LIBS
 565     MLIRCAPIIR
 566     MLIRCAPIPDL
 567 )
 568
 569 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
 570   MODULE_NAME _mlirDialectsSparseTensor
 571   ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
 572   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 573   SOURCES
 574     DialectSparseTensor.cpp
 575   PRIVATE_LINK_LIBS
 576     LLVMSupport
 577   EMBED_CAPI_LINK_LIBS
 578     MLIRCAPIIR
 579     MLIRCAPISparseTensor
 580 )
 581
 582 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
 583   MODULE_NAME _mlirDialectsTransform
 584   ADD_TO_PARENT MLIRPythonSources.Dialects.transform
 585   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 586   SOURCES
 587     DialectTransform.cpp
 588   PRIVATE_LINK_LIBS
 589     LLVMSupport
 590   EMBED_CAPI_LINK_LIBS
 591     MLIRCAPIIR
 592     MLIRCAPITransformDialect
 593 )
 594
 595 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
 596   MODULE_NAME _mlirAsyncPasses
 597   ADD_TO_PARENT MLIRPythonSources.Dialects.async
 598   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 599   SOURCES
 600     AsyncPasses.cpp
 601   PRIVATE_LINK_LIBS
 602     LLVMSupport
 603   EMBED_CAPI_LINK_LIBS
 604     MLIRCAPIAsync
 605 )
 606
 607 if(MLIR_ENABLE_EXECUTION_ENGINE)
 608   declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
 609     MODULE_NAME _mlirExecutionEngine
 610     ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
 611     ROOT_DIR "${PYTHON_SOURCE_DIR}"
 612     SOURCES
 613       ExecutionEngineModule.cpp
 614     PRIVATE_LINK_LIBS
 615       LLVMSupport
 616     EMBED_CAPI_LINK_LIBS
 617       MLIRCAPIExecutionEngine
 618   )
 619 endif()
 620
 621 declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
 622   MODULE_NAME _mlirGPUPasses
 623   ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
 624   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 625   SOURCES
 626     GPUPasses.cpp
 627   PRIVATE_LINK_LIBS
 628     LLVMSupport
 629   EMBED_CAPI_LINK_LIBS
 630     MLIRCAPIGPU
 631 )
 632
 633 declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
 634   MODULE_NAME _mlirLinalgPasses
 635   ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
 636   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 637   SOURCES
 638     LinalgPasses.cpp
 639   PRIVATE_LINK_LIBS
 640     LLVMSupport
 641   EMBED_CAPI_LINK_LIBS
 642     MLIRCAPILinalg
 643 )
 644
 645 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
 646   MODULE_NAME _mlirSparseTensorPasses
 647   ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
 648   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 649   SOURCES
 650     SparseTensorPasses.cpp
 651   PRIVATE_LINK_LIBS
 652     LLVMSupport
 653   EMBED_CAPI_LINK_LIBS
 654     MLIRCAPISparseTensor
 655 )
 656
 657 declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
 658   MODULE_NAME _mlirTransformInterpreter
 659   ADD_TO_PARENT MLIRPythonSources.Dialects.transform
 660   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 661   SOURCES
 662     TransformInterpreter.cpp
 663   PRIVATE_LINK_LIBS
 664     LLVMSupport
 665   EMBED_CAPI_LINK_LIBS
 666     MLIRCAPITransformDialectTransforms
 667 )
 668
 669 # TODO: Figure out how to put this in the test tree.
 670 # This should not be included in the main Python extension. However,
 671 # putting it into MLIRPythonTestSources along with the dialect declaration
 672 # above confuses Python module loader when running under lit.
 673 set(_ADDL_TEST_SOURCES)
 674 if(MLIR_INCLUDE_TESTS)
 675   set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
 676   declare_mlir_python_sources(MLIRPythonTestSources)
 677   declare_mlir_python_sources(MLIRPythonTestSources.Dialects
 678     ADD_TO_PARENT MLIRPythonTestSources)
 679
 680   # TODO: this uses a tablegen file from the test directory and should be
 681   # decoupled from here.
 682   declare_mlir_python_sources(
 683     MLIRPythonTestSources.Dialects.PythonTest
 684     ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 685     ADD_TO_PARENT MLIRPythonTestSources.Dialects
 686     SOURCES dialects/python_test.py)
 687   set(LLVM_TARGET_DEFINITIONS
 688     "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
 689   mlir_tablegen(
 690     "dialects/_python_test_ops_gen.py"
 691     -gen-python-op-bindings
 692     -bind-dialect=python_test)
 693   add_public_tablegen_target(PythonTestDialectPyIncGen)
 694   declare_mlir_python_sources(
 695     MLIRPythonTestSources.Dialects.PythonTest.ops_gen
 696     ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
 697     ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
 698     SOURCES "dialects/_python_test_ops_gen.py")
 699
 700   declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
 701     MODULE_NAME _mlirPythonTest
 702     ADD_TO_PARENT MLIRPythonTestSources.Dialects
 703     ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
 704     SOURCES
 705       PythonTestModule.cpp
 706     PRIVATE_LINK_LIBS
 707       LLVMSupport
 708     EMBED_CAPI_LINK_LIBS
 709       MLIRCAPIPythonTestDialect
 710   )
 711 endif()
 712
 713 ################################################################################
 714 # Common CAPI dependency DSO.
 715 # All python extensions must link through one DSO which exports the CAPI, and
 716 # this must have a globally unique name amongst all embeddors of the python
 717 # library since it will effectively have global scope.
 718 #
 719 # The presence of this aggregate library is part of the long term plan, but its
 720 # use needs to be made more flexible.
 721 #
 722 # TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
 723 # once ready.
 724 ################################################################################
 725
 726 add_mlir_python_common_capi_library(MLIRPythonCAPI
 727   INSTALL_COMPONENT MLIRPythonModules
 728   INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
 729   OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
 730   RELATIVE_INSTALL_ROOT "../../../.."
 731   DECLARED_HEADERS
 732     MLIRPythonCAPI.HeaderSources
 733   DECLARED_SOURCES
 734     MLIRPythonSources
 735     MLIRPythonExtension.RegisterEverything
 736     ${_ADDL_TEST_SOURCES}
 737 )
 738
 739 ################################################################################
 740 # Custom targets.
 741 ################################################################################
 742
 743 _flatten_mlir_python_targets(mlir_python_sources_deps MLIRPythonSources)
 744 add_custom_target("mlir-python-sources" DEPENDS ${mlir_python_sources_deps})
 745 if(NOT LLVM_ENABLE_IDE)
 746   add_llvm_install_targets(install-mlir-python-sources
 747     DEPENDS mlir-python-sources
 748     COMPONENT mlir-python-sources
 749   )
 750 endif()
 751
 752 ################################################################################
 753 # The fully assembled package of modules.
 754 # This must come last.
 755 ################################################################################
 756
 757 add_mlir_python_modules(MLIRPythonModules
 758   ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
 759   INSTALL_PREFIX "python_packages/mlir_core/mlir"
 760   DECLARED_SOURCES
 761     MLIRPythonSources
 762     MLIRPythonExtension.RegisterEverything
 763     ${_ADDL_TEST_SOURCES}
 764   COMMON_CAPI_LINK_LIBS
 765     MLIRPythonCAPI
 766 )