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