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/AsyncOps.td
  53   SOURCES_GLOB dialects/async_dialect/*.py
  54   DIALECT_NAME async_dialect)
  55
  56 declare_mlir_dialect_python_bindings(
  57   ADD_TO_PARENT MLIRPythonSources.Dialects
  58   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
  59   TD_FILE dialects/BufferizationOps.td
  60   SOURCES
  61     dialects/bufferization.py
  62     dialects/_bufferization_ops_ext.py
  63   DIALECT_NAME bufferization)
  64
  65 declare_mlir_dialect_python_bindings(
  66   ADD_TO_PARENT MLIRPythonSources.Dialects
  67   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
  68   TD_FILE dialects/BuiltinOps.td
  69   SOURCES
  70     dialects/builtin.py
  71     dialects/_builtin_ops_ext.py
  72   DIALECT_NAME builtin)
  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/ComplexOps.td
  78   SOURCES
  79     dialects/complex.py
  80   DIALECT_NAME complex)
  81
  82 declare_mlir_dialect_python_bindings(
  83   ADD_TO_PARENT MLIRPythonSources.Dialects
  84   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
  85   TD_FILE dialects/ControlFlowOps.td
  86   SOURCES
  87     dialects/cf.py
  88   DIALECT_NAME cf)
  89
  90 declare_mlir_dialect_python_bindings(
  91   ADD_TO_PARENT MLIRPythonSources.Dialects
  92   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
  93   TD_FILE dialects/FuncOps.td
  94   SOURCES
  95     dialects/func.py
  96     dialects/_func_ops_ext.py
  97   DIALECT_NAME func)
  98
  99 declare_mlir_dialect_python_bindings(
 100   ADD_TO_PARENT MLIRPythonSources.Dialects
 101   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 102   TD_FILE dialects/GPUOps.td
 103   SOURCES_GLOB dialects/gpu/*.py
 104   DIALECT_NAME gpu)
 105
 106 declare_mlir_dialect_python_bindings(
 107   ADD_TO_PARENT MLIRPythonSources.Dialects
 108   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 109   TD_FILE dialects/LinalgOps.td
 110   SOURCES
 111     dialects/_linalg_ops_ext.py
 112   SOURCES_GLOB
 113     dialects/linalg/*.py
 114   DIALECT_NAME linalg
 115   DEPENDS LinalgOdsGen)
 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/TransformOps.td
 121   SOURCES
 122     dialects/_transform_ops_ext.py
 123     dialects/transform/__init__.py
 124     _mlir_libs/_mlir/dialects/transform/__init__.pyi
 125   DIALECT_NAME transform)
 126
 127 declare_mlir_dialect_extension_python_bindings(
 128   ADD_TO_PARENT MLIRPythonSources.Dialects
 129   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 130   TD_FILE dialects/SCFLoopTransformOps.td
 131   SOURCES
 132     dialects/_loop_transform_ops_ext.py
 133     dialects/transform/loop.py
 134   DIALECT_NAME transform
 135   EXTENSION_NAME loop_transform)
 136
 137 declare_mlir_dialect_extension_python_bindings(
 138   ADD_TO_PARENT MLIRPythonSources.Dialects
 139   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 140   TD_FILE dialects/LinalgStructuredTransformOps.td
 141   SOURCES
 142     dialects/_structured_transform_ops_ext.py
 143     dialects/transform/structured.py
 144   DIALECT_NAME transform
 145   EXTENSION_NAME structured_transform)
 146
 147 declare_mlir_dialect_python_bindings(
 148   ADD_TO_PARENT MLIRPythonSources.Dialects
 149   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 150   TD_FILE dialects/MathOps.td
 151   SOURCES dialects/math.py
 152   DIALECT_NAME math)
 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/ArithOps.td
 158   SOURCES
 159     dialects/arith.py
 160     dialects/_arith_ops_ext.py
 161   DIALECT_NAME arith)
 162
 163 declare_mlir_dialect_python_bindings(
 164   ADD_TO_PARENT MLIRPythonSources.Dialects
 165   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 166   TD_FILE dialects/MemRefOps.td
 167   SOURCES
 168     dialects/memref.py
 169     dialects/_memref_ops_ext.py
 170   DIALECT_NAME memref)
 171
 172 declare_mlir_dialect_python_bindings(
 173   ADD_TO_PARENT MLIRPythonSources.Dialects
 174   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 175   TD_FILE dialects/MLProgramOps.td
 176   SOURCES
 177     dialects/ml_program.py
 178     dialects/_ml_program_ops_ext.py
 179   DIALECT_NAME ml_program)
 180
 181 declare_mlir_python_sources(
 182   MLIRPythonSources.Dialects.quant
 183   ADD_TO_PARENT MLIRPythonSources.Dialects
 184   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 185   SOURCES
 186     dialects/quant.py
 187     _mlir_libs/_mlir/dialects/quant.pyi)
 188
 189 declare_mlir_dialect_python_bindings(
 190   ADD_TO_PARENT MLIRPythonSources.Dialects
 191   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 192   TD_FILE dialects/PDLOps.td
 193   SOURCES
 194     dialects/pdl.py
 195     dialects/_pdl_ops_ext.py
 196     _mlir_libs/_mlir/dialects/pdl.pyi
 197   DIALECT_NAME pdl)
 198
 199 declare_mlir_dialect_python_bindings(
 200   ADD_TO_PARENT MLIRPythonSources.Dialects
 201   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 202   TD_FILE dialects/SCFOps.td
 203   SOURCES
 204     dialects/scf.py
 205     dialects/_scf_ops_ext.py
 206   DIALECT_NAME scf)
 207
 208 declare_mlir_dialect_python_bindings(
 209   ADD_TO_PARENT MLIRPythonSources.Dialects
 210   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 211   TD_FILE dialects/ShapeOps.td
 212   SOURCES dialects/shape.py
 213   DIALECT_NAME shape)
 214
 215 declare_mlir_dialect_python_bindings(
 216   ADD_TO_PARENT MLIRPythonSources.Dialects
 217   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 218   TD_FILE dialects/SparseTensorOps.td
 219   SOURCES dialects/sparse_tensor.py
 220   DIALECT_NAME sparse_tensor)
 221
 222 declare_mlir_dialect_python_bindings(
 223   ADD_TO_PARENT MLIRPythonSources.Dialects
 224   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 225   TD_FILE dialects/TensorOps.td
 226   SOURCES
 227     dialects/tensor.py
 228     dialects/_tensor_ops_ext.py
 229   DIALECT_NAME tensor)
 230
 231 declare_mlir_dialect_python_bindings(
 232   ADD_TO_PARENT MLIRPythonSources.Dialects
 233   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 234   TD_FILE dialects/TosaOps.td
 235   SOURCES dialects/tosa.py
 236   DIALECT_NAME tosa)
 237
 238 declare_mlir_dialect_python_bindings(
 239   ADD_TO_PARENT MLIRPythonSources.Dialects
 240   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 241   TD_FILE dialects/VectorOps.td
 242   SOURCES dialects/vector.py
 243   DIALECT_NAME vector)
 244
 245 ################################################################################
 246 # Python extensions.
 247 # The sources for these are all in lib/Bindings/Python, but since they have to
 248 # be rebuilt for each package and integrate with the source setup here, we
 249 # just reference them here instead of having ordered, cross package target
 250 # dependencies.
 251 ################################################################################
 252
 253 set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
 254 declare_mlir_python_extension(MLIRPythonExtension.Core
 255   MODULE_NAME _mlir
 256   ADD_TO_PARENT MLIRPythonSources.Core
 257   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 258   SOURCES
 259     MainModule.cpp
 260     IRAffine.cpp
 261     IRAttributes.cpp
 262     IRCore.cpp
 263     IRInterfaces.cpp
 264     IRModule.cpp
 265     IRTypes.cpp
 266     PybindUtils.cpp
 267     Pass.cpp
 268
 269     # Headers must be included explicitly so they are installed.
 270     Globals.h
 271     IRModule.h
 272     Pass.h
 273     PybindUtils.h
 274   PRIVATE_LINK_LIBS
 275     LLVMSupport
 276   EMBED_CAPI_LINK_LIBS
 277     MLIRCAPIDebug
 278     MLIRCAPIIR
 279     MLIRCAPIInterfaces
 280
 281     # Dialects
 282     MLIRCAPIFunc
 283 )
 284
 285 # This extension exposes an API to register all dialects, extensions, and passes
 286 # packaged in upstream MLIR and it is used for the upstream "mlir" Python
 287 # package. Downstreams will likely want to provide their own and not depend
 288 # on this one, since it links in the world.
 289 # Note that this is not added to any top-level source target for transitive
 290 # inclusion: It must be included explicitly by downstreams if desired. Note that
 291 # this has a very large impact on what gets built/packaged.
 292 declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
 293   MODULE_NAME _mlirRegisterEverything
 294   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 295   SOURCES
 296     RegisterEverything.cpp
 297   PRIVATE_LINK_LIBS
 298     LLVMSupport
 299   EMBED_CAPI_LINK_LIBS
 300     MLIRCAPIConversion
 301     MLIRCAPITransforms
 302     MLIRCAPIRegisterEverything
 303 )
 304
 305 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
 306   MODULE_NAME _mlirDialectsLinalg
 307   ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
 308   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 309   SOURCES
 310     DialectLinalg.cpp
 311   PRIVATE_LINK_LIBS
 312     LLVMSupport
 313   EMBED_CAPI_LINK_LIBS
 314     MLIRCAPIIR
 315     MLIRCAPILinalg
 316 )
 317
 318 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
 319   MODULE_NAME _mlirDialectsQuant
 320   ADD_TO_PARENT MLIRPythonSources.Dialects.quant
 321   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 322   SOURCES
 323     DialectQuant.cpp
 324   PRIVATE_LINK_LIBS
 325     LLVMSupport
 326   EMBED_CAPI_LINK_LIBS
 327     MLIRCAPIIR
 328     MLIRCAPIQuant
 329 )
 330
 331 declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
 332   MODULE_NAME _mlirDialectsPDL
 333   ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
 334   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 335   SOURCES
 336     DialectPDL.cpp
 337   PRIVATE_LINK_LIBS
 338     LLVMSupport
 339   EMBED_CAPI_LINK_LIBS
 340     MLIRCAPIIR
 341     MLIRCAPIPDL
 342 )
 343
 344 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
 345   MODULE_NAME _mlirDialectsSparseTensor
 346   ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
 347   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 348   SOURCES
 349     DialectSparseTensor.cpp
 350   PRIVATE_LINK_LIBS
 351     LLVMSupport
 352   EMBED_CAPI_LINK_LIBS
 353     MLIRCAPIIR
 354     MLIRCAPISparseTensor
 355 )
 356
 357 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
 358   MODULE_NAME _mlirDialectsTransform
 359   ADD_TO_PARENT MLIRPythonSources.Dialects.transform
 360   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 361   SOURCES
 362     DialectTransform.cpp
 363   PRIVATE_LINK_LIBS
 364     LLVMSupport
 365   EMBED_CAPI_LINK_LIBS
 366     MLIRCAPIIR
 367     MLIRCAPITransformDialect
 368 )
 369
 370 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
 371   MODULE_NAME _mlirAsyncPasses
 372   ADD_TO_PARENT MLIRPythonSources.Dialects.async_dialect
 373   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 374   SOURCES
 375     AsyncPasses.cpp
 376   PRIVATE_LINK_LIBS
 377     LLVMSupport
 378   EMBED_CAPI_LINK_LIBS
 379     MLIRCAPIAsync
 380 )
 381
 382 if(MLIR_ENABLE_EXECUTION_ENGINE)
 383   declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine
 384     MODULE_NAME _mlirExecutionEngine
 385     ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
 386     ROOT_DIR "${PYTHON_SOURCE_DIR}"
 387     SOURCES
 388       ExecutionEngineModule.cpp
 389     PRIVATE_LINK_LIBS
 390       LLVMSupport
 391     EMBED_CAPI_LINK_LIBS
 392       MLIRCAPIExecutionEngine
 393   )
 394 endif()
 395
 396 declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
 397   MODULE_NAME _mlirGPUPasses
 398   ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
 399   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 400   SOURCES
 401     GPUPasses.cpp
 402   PRIVATE_LINK_LIBS
 403     LLVMSupport
 404   EMBED_CAPI_LINK_LIBS
 405     MLIRCAPIGPU
 406 )
 407
 408 declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
 409   MODULE_NAME _mlirLinalgPasses
 410   ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
 411   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 412   SOURCES
 413     LinalgPasses.cpp
 414   PRIVATE_LINK_LIBS
 415     LLVMSupport
 416   EMBED_CAPI_LINK_LIBS
 417     MLIRCAPILinalg
 418 )
 419
 420 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
 421   MODULE_NAME _mlirSparseTensorPasses
 422   ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
 423   ROOT_DIR "${PYTHON_SOURCE_DIR}"
 424   SOURCES
 425     SparseTensorPasses.cpp
 426   PRIVATE_LINK_LIBS
 427     LLVMSupport
 428   EMBED_CAPI_LINK_LIBS
 429     MLIRCAPISparseTensor
 430 )
 431
 432 # TODO: Figure out how to put this in the test tree.
 433 # This should not be included in the main Python extension. However,
 434 # putting it into MLIRPythonTestSources along with the dialect declaration
 435 # above confuses Python module loader when running under lit.
 436 set(_ADDL_TEST_SOURCES)
 437 if(MLIR_INCLUDE_TESTS)
 438   set(_ADDL_TEST_SOURCES MLIRPythonTestSources)
 439   declare_mlir_python_sources(MLIRPythonTestSources)
 440   declare_mlir_python_sources(MLIRPythonTestSources.Dialects
 441     ADD_TO_PARENT MLIRPythonTestSources)
 442
 443   # TODO: this uses a tablegen file from the test directory and should be
 444   # decoupled from here.
 445   declare_mlir_python_sources(
 446     MLIRPythonTestSources.Dialects.PythonTest
 447     ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
 448     ADD_TO_PARENT MLIRPythonTestSources.Dialects
 449     SOURCES dialects/python_test.py)
 450   set(LLVM_TARGET_DEFINITIONS
 451     "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td")
 452   mlir_tablegen(
 453     "dialects/_python_test_ops_gen.py"
 454     -gen-python-op-bindings
 455     -bind-dialect=python_test)
 456   add_public_tablegen_target(PythonTestDialectPyIncGen)
 457   declare_mlir_python_sources(
 458     MLIRPythonTestSources.Dialects.PythonTest.ops_gen
 459     ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}"
 460     ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest
 461     SOURCES "dialects/_python_test_ops_gen.py")
 462
 463   declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtension
 464     MODULE_NAME _mlirPythonTest
 465     ADD_TO_PARENT MLIRPythonTestSources.Dialects
 466     ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
 467     SOURCES
 468       PythonTestModule.cpp
 469     PRIVATE_LINK_LIBS
 470       LLVMSupport
 471     EMBED_CAPI_LINK_LIBS
 472       MLIRCAPIPythonTestDialect
 473   )
 474 endif()
 475
 476 ################################################################################
 477 # Common CAPI dependency DSO.
 478 # All python extensions must link through one DSO which exports the CAPI, and
 479 # this must have a globally unique name amongst all embeddors of the python
 480 # library since it will effectively have global scope.
 481 #
 482 # The presence of this aggregate library is part of the long term plan, but its
 483 # use needs to be made more flexible.
 484 #
 485 # TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419
 486 # once ready.
 487 ################################################################################
 488
 489 add_mlir_python_common_capi_library(MLIRPythonCAPI
 490   INSTALL_COMPONENT MLIRPythonModules
 491   INSTALL_DESTINATION python_packages/mlir_core/mlir/_mlir_libs
 492   OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
 493   RELATIVE_INSTALL_ROOT "../../../.."
 494   DECLARED_HEADERS
 495     MLIRPythonCAPI.HeaderSources
 496   DECLARED_SOURCES
 497     MLIRPythonSources
 498     MLIRPythonExtension.RegisterEverything
 499     ${_ADDL_TEST_SOURCES}
 500 )
 501
 502 ################################################################################
 503 # The fully assembled package of modules.
 504 # This must come last.
 505 ################################################################################
 506
 507 add_mlir_python_modules(MLIRPythonModules
 508   ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
 509   INSTALL_PREFIX "python_packages/mlir_core/mlir"
 510   DECLARED_SOURCES
 511     MLIRPythonSources
 512     MLIRPythonExtension.RegisterEverything
 513     ${_ADDL_TEST_SOURCES}
 514   COMMON_CAPI_LINK_LIBS
 515     MLIRPythonCAPI
 516 )