11include (AddMLIRPython)
22
3- # Specifies that all MLIR packages are co-located under the `mlir_standalone`
4- # top level package (the API has been embedded in a relocatable way).
5- add_compile_definitions ("MLIR_PYTHON_PACKAGE_PREFIX=${MLIR_PYTHON_PACKAGE_PREFIX} ." )
6-
73################################################################################
84# Structural groupings.
95################################################################################
@@ -27,6 +23,11 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python
2723 passmanager.py
2824 rewrite.py
2925 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
3031)
3132
3233declare_mlir_python_sources(MLIRPythonSources.Core.Python.Extras
@@ -42,6 +43,7 @@ declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
4243 ADD_TO_PARENT MLIRPythonSources
4344 SOURCES
4445 execution_engine.py
46+ _mlir_libs/_mlirExecutionEngine.pyi
4547 SOURCES_GLOB
4648 runtime/*.py
4749)
@@ -193,6 +195,7 @@ declare_mlir_dialect_python_bindings(
193195 TD_FILE dialects/TransformOps.td
194196 SOURCES
195197 dialects/transform/__init__.py
198+ _mlir_libs/_mlir/dialects/transform/__init__.pyi
196199 DIALECT_NAME transform
197200 GEN_ENUM_BINDINGS_TD_FILE
198201 "../../include/mlir/Dialect/Transform/IR/TransformAttrs.td"
@@ -364,7 +367,8 @@ declare_mlir_python_sources(
364367 ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR} /mlir"
365368 GEN_ENUM_BINDINGS
366369 SOURCES
367- dialects/quant.py)
370+ dialects/quant.py
371+ _mlir_libs/_mlir/dialects/quant.pyi)
368372
369373declare_mlir_dialect_python_bindings(
370374 ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -380,6 +384,7 @@ declare_mlir_dialect_python_bindings(
380384 TD_FILE dialects/PDLOps.td
381385 SOURCES
382386 dialects/pdl.py
387+ _mlir_libs/_mlir/dialects/pdl.pyi
383388 DIALECT_NAME pdl)
384389
385390declare_mlir_dialect_python_bindings(
@@ -505,11 +510,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
505510
506511 # Dialects
507512 MLIRCAPIFunc
508- GENERATE_TYPE_STUBS
509- "_mlir/__init__.pyi"
510- "_mlir/ir.pyi"
511- "_mlir/passmanager.pyi"
512- "_mlir/rewrite.pyi"
513513)
514514
515515# This extension exposes an API to register all dialects, extensions, and passes
@@ -531,8 +531,6 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
531531 MLIRCAPIConversion
532532 MLIRCAPITransforms
533533 MLIRCAPIRegisterEverything
534- GENERATE_TYPE_STUBS
535- "_mlirRegisterEverything.pyi"
536534)
537535
538536declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -547,8 +545,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
547545 EMBED_CAPI_LINK_LIBS
548546 MLIRCAPIIR
549547 MLIRCAPILinalg
550- GENERATE_TYPE_STUBS
551- "_mlirDialectsLinalg.pyi"
552548)
553549
554550declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -563,8 +559,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
563559 EMBED_CAPI_LINK_LIBS
564560 MLIRCAPIIR
565561 MLIRCAPIGPU
566- GENERATE_TYPE_STUBS
567- "_mlirDialectsGPU.pyi"
568562)
569563
570564declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -579,8 +573,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
579573 EMBED_CAPI_LINK_LIBS
580574 MLIRCAPIIR
581575 MLIRCAPILLVM
582- GENERATE_TYPE_STUBS
583- "_mlirDialectsLLVM.pyi"
584576)
585577
586578declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -595,8 +587,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
595587 EMBED_CAPI_LINK_LIBS
596588 MLIRCAPIIR
597589 MLIRCAPIQuant
598- GENERATE_TYPE_STUBS
599- "_mlirDialectsQuant.pyi"
600590)
601591
602592declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -611,8 +601,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
611601 EMBED_CAPI_LINK_LIBS
612602 MLIRCAPIIR
613603 MLIRCAPINVGPU
614- GENERATE_TYPE_STUBS
615- "_mlirDialectsNVGPU.pyi"
616604)
617605
618606declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -627,8 +615,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
627615 EMBED_CAPI_LINK_LIBS
628616 MLIRCAPIIR
629617 MLIRCAPIPDL
630- GENERATE_TYPE_STUBS
631- "_mlirDialectsPDL.pyi"
632618)
633619
634620declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -643,8 +629,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
643629 EMBED_CAPI_LINK_LIBS
644630 MLIRCAPIIR
645631 MLIRCAPISparseTensor
646- GENERATE_TYPE_STUBS
647- "_mlirDialectsSparseTensor.pyi"
648632)
649633
650634declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -659,8 +643,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
659643 EMBED_CAPI_LINK_LIBS
660644 MLIRCAPIIR
661645 MLIRCAPITransformDialect
662- GENERATE_TYPE_STUBS
663- "_mlirDialectsTransform.pyi"
664646)
665647
666648declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -674,8 +656,6 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
674656 LLVMSupport
675657 EMBED_CAPI_LINK_LIBS
676658 MLIRCAPIAsync
677- GENERATE_TYPE_STUBS
678- "_mlirAsyncPasses.pyi"
679659)
680660
681661if (MLIR_ENABLE_EXECUTION_ENGINE)
@@ -690,8 +670,6 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
690670 LLVMSupport
691671 EMBED_CAPI_LINK_LIBS
692672 MLIRCAPIExecutionEngine
693- GENERATE_TYPE_STUBS
694- "_mlirExecutionEngine.pyi"
695673 )
696674endif ()
697675
@@ -706,8 +684,6 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
706684 LLVMSupport
707685 EMBED_CAPI_LINK_LIBS
708686 MLIRCAPIGPU
709- GENERATE_TYPE_STUBS
710- "_mlirGPUPasses.pyi"
711687)
712688
713689declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -721,8 +697,6 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
721697 LLVMSupport
722698 EMBED_CAPI_LINK_LIBS
723699 MLIRCAPILinalg
724- GENERATE_TYPE_STUBS
725- "_mlirLinalgPasses.pyi"
726700)
727701
728702declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -740,8 +714,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
740714 MLIRCAPIIR
741715 MLIRCAPISMT
742716 MLIRCAPIExportSMTLIB
743- GENERATE_TYPE_STUBS
744- "_mlirDialectsSMT.pyi"
745717)
746718
747719declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -755,8 +727,6 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
755727 LLVMSupport
756728 EMBED_CAPI_LINK_LIBS
757729 MLIRCAPISparseTensor
758- GENERATE_TYPE_STUBS
759- "_mlirSparseTensorPasses.pyi"
760730)
761731
762732declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -770,8 +740,6 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
770740 LLVMSupport
771741 EMBED_CAPI_LINK_LIBS
772742 MLIRCAPITransformDialectTransforms
773- GENERATE_TYPE_STUBS
774- "_mlirTransformInterpreter.pyi"
775743)
776744
777745# TODO: Figure out how to put this in the test tree.
@@ -830,8 +798,6 @@ if(MLIR_INCLUDE_TESTS)
830798 LLVMSupport
831799 EMBED_CAPI_LINK_LIBS
832800 MLIRCAPIPythonTestDialect
833- GENERATE_TYPE_STUBS
834- "_mlirPythonTestNanobind.pyi"
835801 )
836802endif ()
837803
@@ -851,7 +817,7 @@ endif()
851817add_mlir_python_common_capi_library(MLIRPythonCAPI
852818 INSTALL_COMPONENT MLIRPythonModules
853819 INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX} /_mlir_libs"
854- OUTPUT_DIRECTORY "${MLIR_BINARY_DIR} /${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX} /_mlir_libs"
820+ OUTPUT_DIRECTORY "${MLIR_BINARY_DIR} /python_packages/mlir_core/mlir /_mlir_libs"
855821 RELATIVE_INSTALL_ROOT "../../../.."
856822 DECLARED_HEADERS
857823 MLIRPythonCAPI.HeaderSources
@@ -880,7 +846,7 @@ endif()
880846################################################################################
881847
882848add_mlir_python_modules(MLIRPythonModules
883- ROOT_PREFIX "${MLIR_BINARY_DIR} /${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX} "
849+ ROOT_PREFIX "${MLIR_BINARY_DIR} /python_packages/mlir_core/mlir "
884850 INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX} "
885851 DECLARED_SOURCES
886852 MLIRPythonSources
0 commit comments