include(AddMLIRPython) # Specifies that all MLIR packages are co-located under the `MLIR_PYTHON_PACKAGE_PREFIX.` # top level package (the API has been embedded in a relocatable way). add_compile_definitions("MLIR_PYTHON_PACKAGE_PREFIX=${MLIR_PYTHON_PACKAGE_PREFIX}.") set(MLIRPythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}") set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python") ################################################################################ # Structural groupings. ################################################################################ declare_mlir_python_sources(MLIRPythonSources) declare_mlir_python_sources(MLIRPythonSources.Dialects ADD_TO_PARENT MLIRPythonSources) declare_mlir_python_sources(MLIRPythonSources.Core ADD_TO_PARENT MLIRPythonSources) ################################################################################ # Pure python sources and generated code ################################################################################ declare_mlir_python_sources(MLIRPythonSources.Core.Python ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" ADD_TO_PARENT MLIRPythonSources.Core SOURCES _mlir_libs/__init__.py _mlir_libs/_mlir/py.typed ir.py passmanager.py rewrite.py dialects/_ods_common.py dialects/ext.py ) declare_mlir_python_sources(MLIRPythonSources.Core.Python.Extras ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" ADD_TO_PARENT MLIRPythonSources.Core.Python SOURCES extras/types.py extras/meta.py ) declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" ADD_TO_PARENT MLIRPythonSources SOURCES execution_engine.py _mlir_libs/_mlirExecutionEngine.pyi SOURCES_GLOB runtime/*.py ) declare_mlir_python_sources(MLIRPythonCAPI.HeaderSources ROOT_DIR "${MLIR_SOURCE_DIR}/include" SOURCES_GLOB "mlir-c/*.h" ) ################################################################################ # Dialect bindings ################################################################################ declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/AffineOps.td SOURCES dialects/affine.py DIALECT_NAME affine GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/AMDGPUOps.td SOURCES dialects/amdgpu.py DIALECT_NAME amdgpu GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/AsyncOps.td SOURCES_GLOB dialects/async_dialect/*.py DIALECT_NAME async) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/BufferizationOps.td SOURCES dialects/bufferization.py DIALECT_NAME bufferization GEN_ENUM_BINDINGS_TD_FILE "../../include/mlir/Dialect/Bufferization/IR/BufferizationEnums.td" ) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/BuiltinOps.td SOURCES dialects/builtin.py DIALECT_NAME builtin) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/ComplexOps.td SOURCES dialects/complex.py DIALECT_NAME complex) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/IndexOps.td SOURCES dialects/index.py DIALECT_NAME index GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/ControlFlowOps.td SOURCES dialects/cf.py DIALECT_NAME cf) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/FuncOps.td SOURCES dialects/func.py DIALECT_NAME func) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/OpenACCOps.td SOURCES dialects/openacc.py DIALECT_NAME acc DEPENDS acc_common_td ) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/GPUOps.td SOURCES_GLOB dialects/gpu/*.py DIALECT_NAME gpu GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/LinalgOps.td SOURCES SOURCES_GLOB dialects/linalg/*.py DIALECT_NAME linalg DEPENDS LinalgOdsGen GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/LLVMOps.td SOURCES dialects/llvm.py DIALECT_NAME llvm GEN_ENUM_BINDINGS) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/TransformPDLExtensionOps.td SOURCES dialects/transform/pdl.py DIALECT_NAME transform EXTENSION_NAME transform_pdl_extension) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/TransformSMTExtensionOps.td SOURCES dialects/transform/smt.py DIALECT_NAME transform EXTENSION_NAME transform_smt_extension) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/TransformDebugExtensionOps.td SOURCES dialects/transform/debug.py DIALECT_NAME transform EXTENSION_NAME transform_debug_extension) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/TransformTuneExtensionOps.td SOURCES dialects/transform/tune.py DIALECT_NAME transform EXTENSION_NAME transform_tune_extension) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/TransformOps.td SOURCES dialects/transform/__init__.py _mlir_libs/_mlir/dialects/transform/__init__.pyi DIALECT_NAME transform GEN_ENUM_BINDINGS_TD_FILE "../../include/mlir/Dialect/Transform/IR/TransformAttrs.td" ) declare_mlir_python_sources( MLIRPythonSources.Dialects.transform.extras ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" GEN_ENUM_BINDINGS SOURCES dialects/transform/extras/__init__.py) declare_mlir_python_sources( MLIRPythonSources.Dialects.transform.interpreter ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" SOURCES dialects/transform/interpreter/__init__.py) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/BufferizationTransformOps.td SOURCES dialects/transform/bufferization.py DIALECT_NAME transform EXTENSION_NAME bufferization_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/GPUTransformOps.td SOURCES dialects/transform/gpu.py DIALECT_NAME transform EXTENSION_NAME gpu_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/SCFLoopTransformOps.td SOURCES dialects/transform/loop.py DIALECT_NAME transform EXTENSION_NAME loop_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/MemRefTransformOps.td SOURCES dialects/transform/memref.py DIALECT_NAME transform EXTENSION_NAME memref_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/NVGPUTransformOps.td SOURCES dialects/transform/nvgpu.py DIALECT_NAME transform EXTENSION_NAME nvgpu_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/LinalgStructuredTransformOps.td SOURCES dialects/transform/structured.py DIALECT_NAME transform EXTENSION_NAME structured_transform GEN_ENUM_BINDINGS_TD_FILE "../../include/mlir/Dialect/Linalg/TransformOps/LinalgTransformEnums.td" ) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/SparseTensorTransformOps.td SOURCES dialects/transform/sparse_tensor.py DIALECT_NAME transform EXTENSION_NAME sparse_tensor_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/TensorTransformOps.td SOURCES dialects/transform/tensor.py DIALECT_NAME transform EXTENSION_NAME tensor_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/VectorTransformOps.td SOURCES dialects/transform/vector.py DIALECT_NAME transform EXTENSION_NAME vector_transform GEN_ENUM_BINDINGS_TD_FILE "../../include/mlir/Dialect/Vector/Transforms/VectorTransformsBase.td" ) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/X86TransformOps.td SOURCES dialects/transform/x86.py DIALECT_NAME transform EXTENSION_NAME x86_transform) declare_mlir_dialect_extension_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/XeGPUTransformOps.td SOURCES dialects/transform/xegpu.py DIALECT_NAME transform EXTENSION_NAME xegpu_transform) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/MathOps.td SOURCES dialects/math.py DIALECT_NAME math) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/ArithOps.td SOURCES dialects/arith.py DIALECT_NAME arith GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/MemRefOps.td SOURCES dialects/memref.py DIALECT_NAME memref) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/ShardOps.td SOURCES dialects/shard.py DIALECT_NAME shard GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/MLProgramOps.td SOURCES dialects/ml_program.py DIALECT_NAME ml_program) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/NVGPUOps.td SOURCES dialects/nvgpu.py DIALECT_NAME nvgpu GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/NVVMOps.td SOURCES dialects/nvvm.py DIALECT_NAME nvvm GEN_ENUM_BINDINGS) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/ROCDLOps.td SOURCES dialects/rocdl.py DIALECT_NAME rocdl) declare_mlir_python_sources( MLIRPythonSources.Dialects.quant ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" GEN_ENUM_BINDINGS SOURCES dialects/quant.py _mlir_libs/_mlir/dialects/quant.pyi) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/EmitC.td SOURCES dialects/emitc.py DIALECT_NAME emitc) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/PDLOps.td SOURCES dialects/pdl.py _mlir_libs/_mlir/dialects/pdl.pyi DIALECT_NAME pdl) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/OpenMPOps.td SOURCES dialects/openmp.py DIALECT_NAME omp DEPENDS omp_common_td) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/SCFOps.td SOURCES dialects/scf.py DIALECT_NAME scf) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/ShapeOps.td SOURCES dialects/shape.py DIALECT_NAME shape) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/SparseTensorOps.td SOURCES dialects/sparse_tensor.py DIALECT_NAME sparse_tensor GEN_ENUM_BINDINGS_TD_FILE "../../include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td" ) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/SMTOps.td GEN_ENUM_BINDINGS SOURCES dialects/smt.py DIALECT_NAME smt) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/SPIRVOps.td SOURCES dialects/spirv.py DIALECT_NAME spirv) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/TensorOps.td SOURCES dialects/tensor.py DIALECT_NAME tensor) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/TosaOps.td SOURCES dialects/tosa.py DIALECT_NAME tosa ) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/UBOps.td SOURCES dialects/ub.py DIALECT_NAME ub ) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/VectorOps.td SOURCES dialects/vector.py DIALECT_NAME vector GEN_ENUM_BINDINGS_TD_FILE "dialects/VectorAttributes.td") declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/X86.td SOURCES dialects/x86.py DIALECT_NAME x86) declare_mlir_dialect_python_bindings( ADD_TO_PARENT MLIRPythonSources.Dialects ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" TD_FILE dialects/IRDLOps.td SOURCES dialects/irdl.py DIALECT_NAME irdl GEN_ENUM_BINDINGS ) ################################################################################ # Python extensions. # The sources for these are all in lib/Bindings/Python, but since they have to # be rebuilt for each package and integrate with the source setup here, we # just reference them here instead of having ordered, cross package target # dependencies. ################################################################################ declare_mlir_python_extension(MLIRPythonExtension.Core MODULE_NAME _mlir ADD_TO_PARENT MLIRPythonSources.Core ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES MainModule.cpp IRAffine.cpp IRInterfaces.cpp Pass.cpp Rewrite.cpp # Headers must be included explicitly so they are installed. Pass.h Rewrite.h PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIDebug MLIRCAPIIR MLIRCAPIInterfaces # Dialects MLIRCAPIFunc ) # This extension exposes an API to register all dialects, extensions, and passes # packaged in upstream MLIR and it is used for the upstream "mlir" Python # package. Downstreams will likely want to provide their own and not depend # on this one, since it links in the world. # Note that this is not added to any top-level source target for transitive # inclusion: It must be included explicitly by downstreams if desired. Note that # this has a very large impact on what gets built/packaged. declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything MODULE_NAME _mlirRegisterEverything ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES RegisterEverything.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIConversion MLIRCAPITransforms MLIRCAPIRegisterEverything ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Nanobind MODULE_NAME _mlirDialectsLinalg ADD_TO_PARENT MLIRPythonSources.Dialects.linalg ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectLinalg.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPILinalg ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Nanobind MODULE_NAME _mlirDialectsGPU ADD_TO_PARENT MLIRPythonSources.Dialects.gpu ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectGPU.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPIGPU ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Nanobind MODULE_NAME _mlirDialectsLLVM ADD_TO_PARENT MLIRPythonSources.Dialects.llvm ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectLLVM.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPILLVM # Misnomer - this is only the LLVMIR translation target. MLIRCAPITarget ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Nanobind MODULE_NAME _mlirDialectsQuant ADD_TO_PARENT MLIRPythonSources.Dialects.quant ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectQuant.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPIQuant ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Nanobind MODULE_NAME _mlirDialectsNVGPU ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectNVGPU.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPINVGPU ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Nanobind MODULE_NAME _mlirDialectsPDL ADD_TO_PARENT MLIRPythonSources.Dialects.pdl ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectPDL.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPIPDL ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Nanobind MODULE_NAME _mlirDialectsSparseTensor ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectSparseTensor.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPISparseTensor ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Nanobind MODULE_NAME _mlirDialectsTransform ADD_TO_PARENT MLIRPythonSources.Dialects.transform ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectTransform.cpp Rewrite.h PRIVATE_LINK_LIBS LLVMSupport MLIRPythonExtension.Core EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPITransformDialect ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.IRDL.Nanobind MODULE_NAME _mlirDialectsIRDL ADD_TO_PARENT MLIRPythonSources.Dialects.irdl ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectIRDL.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPIIRDL ) declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses MODULE_NAME _mlirAsyncPasses ADD_TO_PARENT MLIRPythonSources.Dialects.async ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES AsyncPasses.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIAsync ) if(MLIR_ENABLE_EXECUTION_ENGINE) declare_mlir_python_extension(MLIRPythonExtension.ExecutionEngine MODULE_NAME _mlirExecutionEngine ADD_TO_PARENT MLIRPythonSources.ExecutionEngine ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES ExecutionEngineModule.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIExecutionEngine ) endif() declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses MODULE_NAME _mlirGPUPasses ADD_TO_PARENT MLIRPythonSources.Dialects.gpu ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES GPUPasses.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIGPU ) declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses MODULE_NAME _mlirLinalgPasses ADD_TO_PARENT MLIRPythonSources.Dialects.linalg ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES LinalgPasses.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPILinalg ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Nanobind MODULE_NAME _mlirDialectsSMT ADD_TO_PARENT MLIRPythonSources.Dialects.smt ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectSMT.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPISMT MLIRCAPIExportSMTLIB ) declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses MODULE_NAME _mlirSparseTensorPasses ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES SparseTensorPasses.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPISparseTensor ) declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter MODULE_NAME _mlirTransformInterpreter ADD_TO_PARENT MLIRPythonSources.Dialects.transform ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES TransformInterpreter.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPITransformDialectTransforms ) declare_mlir_python_extension(MLIRPythonExtension.Dialects.AMDGPU.Nanobind MODULE_NAME _mlirDialectsAMDGPU ADD_TO_PARENT MLIRPythonSources.Dialects.amdgpu ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES DialectAMDGPU.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIIR MLIRCAPIAMDGPU ) # TODO: Figure out how to put this in the test tree. # This should not be included in the main Python extension. However, # putting it into MLIRPythonTestSources along with the dialect declaration # above confuses Python module loader when running under lit. set(_ADDL_TEST_SOURCES) if(MLIR_INCLUDE_TESTS) set(_ADDL_TEST_SOURCES MLIRPythonTestSources) declare_mlir_python_sources(MLIRPythonTestSources) declare_mlir_python_sources(MLIRPythonTestSources.Dialects ADD_TO_PARENT MLIRPythonTestSources) # TODO: this uses a tablegen file from the test directory and should be # decoupled from here. declare_mlir_python_sources( MLIRPythonTestSources.Dialects.PythonTest ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir" ADD_TO_PARENT MLIRPythonTestSources.Dialects SOURCES dialects/python_test.py ) set(LLVM_TARGET_DEFINITIONS "${MLIR_MAIN_SRC_DIR}/test/python/python_test_ops.td") mlir_tablegen( "dialects/_python_test_ops_gen.py" -gen-python-op-bindings -bind-dialect=python_test) add_public_tablegen_target(PythonTestDialectPyIncGen) declare_mlir_python_sources( MLIRPythonTestSources.Dialects.PythonTest.ops_gen ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}" ADD_TO_PARENT MLIRPythonTestSources.Dialects.PythonTest SOURCES "dialects/_python_test_ops_gen.py") declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionNanobind MODULE_NAME _mlirPythonTestNanobind ADD_TO_PARENT MLIRPythonTestSources.Dialects ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib" SOURCES PythonTestModuleNanobind.cpp PRIVATE_LINK_LIBS LLVMSupport EMBED_CAPI_LINK_LIBS MLIRCAPIPythonTestDialect ) endif() declare_mlir_python_extension(MLIRPythonExtension.MLIRPythonSupport _PRIVATE_SUPPORT_LIB MODULE_NAME MLIRPythonSupport ADD_TO_PARENT MLIRPythonSources.Core ROOT_DIR "${PYTHON_SOURCE_DIR}" SOURCES Globals.cpp IRAttributes.cpp IRCore.cpp IRTypes.cpp ) ################################################################################ # Common CAPI dependency DSO. # All python extensions must link through one DSO which exports the CAPI, and # this must have a globally unique name amongst all embeddors of the python # library since it will effectively have global scope. # # The presence of this aggregate library is part of the long term plan, but its # use needs to be made more flexible. # # TODO: Upgrade to the aggregate utility in https://reviews.llvm.org/D106419 # once ready. ################################################################################ add_mlir_python_common_capi_library(MLIRPythonCAPI INSTALL_COMPONENT MLIRPythonModules INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs" OUTPUT_DIRECTORY "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs" RELATIVE_INSTALL_ROOT "../../../.." DECLARED_HEADERS MLIRPythonCAPI.HeaderSources DECLARED_SOURCES MLIRPythonSources MLIRPythonExtension.RegisterEverything ${_ADDL_TEST_SOURCES} ) ################################################################################ # Custom targets. ################################################################################ _flatten_mlir_python_targets(mlir_python_sources_deps MLIRPythonSources) if(MLIR_PYTHON_STUBGEN_ENABLED) # _mlir stubgen # Note: All this needs to come before add_mlir_python_modules(MLIRPythonModules so that the install targets for the # generated type stubs get created. set(_core_type_stub_sources _mlir/__init__.pyi _mlir/ir.pyi _mlir/passmanager.pyi _mlir/rewrite.pyi ) # Note 1: INTERFACE_SOURCES is a genex ($ $) # which will be evaluated by file(GENERATE ...) inside mlir_generate_type_stubs. This will evaluate to the correct # thing in the build dir (i.e., actual source dir paths) and in the install dir # (where it's a conventional path; see install/lib/cmake/mlir/MLIRTargets.cmake). # # Note 2: MLIRPythonExtension.Core is the target that is defined using target_sources(INTERFACE) # **NOT** MLIRPythonModules.extension._mlir.dso. So be sure to use the correct target! get_target_property(_core_extension_srcs MLIRPythonExtension.Core INTERFACE_SOURCES) # Why is MODULE_NAME _mlir here but mlir._mlir_libs._mlirPythonTestNanobind below??? # The _mlir extension can be imported independently of any other python code and/or extension modules. # I.e., you could do `cd $MLIRPythonModules_ROOT_PREFIX/_mlir_libs && python -c "import _mlir"` (try it!). # _mlir is also (currently) the only extension for which this is possible because dialect extensions modules, # which generally make use of `mlir_value_subclass/mlir_type_subclass/mlir_attribute_subclass`, perform an # `import mlir` right when they're loaded (see the mlir_*_subclass ctors in NanobindAdaptors.h). # Note, this also why IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs" here while below # "${MLIRPythonModules_ROOT_PREFIX}/.." (because MLIR_BINDINGS_PYTHON_INSTALL_PREFIX, by default, ends at mlir). # # Further note: this function creates file targets like # "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs/_mlir/__init__.pyi". These must match the file targets # that declare_mlir_python_sources expects, which are like "${ROOT_DIR}/${WHATEVER_SOURCE}". # This is why _mlir_libs is prepended below. mlir_generate_type_stubs( MODULE_NAME _mlir DEPENDS_TARGETS MLIRPythonModules.extension._mlir.dso OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" OUTPUTS "${_core_type_stub_sources}" DEPENDS_TARGET_SRC_DEPS "${_core_extension_srcs}" IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/_mlir_libs" ) set(_mlir_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}") list(TRANSFORM _core_type_stub_sources PREPEND "_mlir_libs/") # Note, we do not do ADD_TO_PARENT here so that the type stubs are not associated (as mlir_DEPENDS) with # MLIRPythonSources.Core (or something) when a distro is installed/created. Otherwise they would not be regenerated # by users of the distro (the stubs are still installed in the distro - they are just not added to mlir_DEPENDS). declare_mlir_python_sources( MLIRPythonExtension.Core.type_stub_gen ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" SOURCES "${_core_type_stub_sources}" ) list(APPEND mlir_python_sources_deps MLIRPythonExtension.Core.type_stub_gen) # _mlirPythonTestNanobind stubgen if(MLIR_INCLUDE_TESTS) get_target_property(_test_extension_srcs MLIRPythonTestSources.PythonTestExtensionNanobind INTERFACE_SOURCES) mlir_generate_type_stubs( # This is the FQN path because dialect modules import _mlir when loaded. See above. MODULE_NAME mlir._mlir_libs._mlirPythonTestNanobind DEPENDS_TARGETS # You need both _mlir and _mlirPythonTestNanobind because dialect modules import _mlir when loaded # (so _mlir needs to be built before calling stubgen). MLIRPythonModules.extension._mlir.dso MLIRPythonModules.extension._mlirPythonTestNanobind.dso # You need this one so that ir.py "built" because mlir._mlir_libs.__init__.py import mlir.ir in _site_initialize. MLIRPythonModules.sources.MLIRPythonSources.Core.Python OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs" OUTPUTS _mlirPythonTestNanobind.pyi DEPENDS_TARGET_SRC_DEPS "${_test_extension_srcs}" IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}/.." ) set(_mlirPythonTestNanobind_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}") declare_mlir_python_sources( MLIRPythonTestSources.PythonTestExtensionNanobind.type_stub_gen ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs" ADD_TO_PARENT MLIRPythonTestSources.Dialects SOURCES _mlir_libs/_mlirPythonTestNanobind.pyi ) endif() endif() add_custom_target("mlir-python-sources" DEPENDS ${mlir_python_sources_deps}) if(NOT LLVM_ENABLE_IDE) add_llvm_install_targets(install-mlir-python-sources DEPENDS mlir-python-sources COMPONENT mlir-python-sources ) endif() ################################################################################ # The fully assembled package of modules. # This must come last. ################################################################################ set(_declared_sources MLIRPythonSources MLIRPythonExtension.RegisterEverything) if(MLIR_PYTHON_STUBGEN_ENABLED) list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen) endif() add_mlir_python_modules(MLIRPythonModules ROOT_PREFIX ${MLIRPythonModules_ROOT_PREFIX} INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}" DECLARED_SOURCES ${_declared_sources} ${_ADDL_TEST_SOURCES} COMMON_CAPI_LINK_LIBS MLIRPythonCAPI ) if(MLIR_PYTHON_STUBGEN_ENABLED) add_dependencies(MLIRPythonModules "${_mlir_typestub_gen_target}") if(MLIR_INCLUDE_TESTS) add_dependencies(MLIRPythonModules "${_mlirPythonTestNanobind_typestub_gen_target}") endif() endif()