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}.")

################################################################################
# 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
)

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/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/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.
################################################################################

set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
declare_mlir_python_extension(MLIRPythonExtension.Core
  MODULE_NAME _mlir
  ADD_TO_PARENT MLIRPythonSources.Core
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    MainModule.cpp
    IRAffine.cpp
    IRAttributes.cpp
    IRCore.cpp
    IRInterfaces.cpp
    IRModule.cpp
    IRTypes.cpp
    Pass.cpp
    Rewrite.cpp

    # Headers must be included explicitly so they are installed.
    Globals.h
    IRModule.h
    Pass.h
    NanobindUtils.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}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    RegisterEverything.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPIConversion
    MLIRCAPITransforms
    MLIRCAPIRegisterEverything
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
  MODULE_NAME _mlirDialectsLinalg
  ADD_TO_PARENT MLIRPythonSources.Dialects.linalg
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectLinalg.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPIIR
    MLIRCAPILinalg
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
  MODULE_NAME _mlirDialectsGPU
  ADD_TO_PARENT MLIRPythonSources.Dialects.gpu
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectGPU.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPIIR
    MLIRCAPIGPU
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
  MODULE_NAME _mlirDialectsLLVM
  ADD_TO_PARENT MLIRPythonSources.Dialects.llvm
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  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.Pybind
  MODULE_NAME _mlirDialectsQuant
  ADD_TO_PARENT MLIRPythonSources.Dialects.quant
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectQuant.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPIIR
    MLIRCAPIQuant
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
  MODULE_NAME _mlirDialectsNVGPU
  ADD_TO_PARENT MLIRPythonSources.Dialects.nvgpu
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectNVGPU.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPIIR
    MLIRCAPINVGPU
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
  MODULE_NAME _mlirDialectsPDL
  ADD_TO_PARENT MLIRPythonSources.Dialects.pdl
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectPDL.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPIIR
    MLIRCAPIPDL
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
  MODULE_NAME _mlirDialectsSparseTensor
  ADD_TO_PARENT MLIRPythonSources.Dialects.sparse_tensor
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectSparseTensor.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPIIR
    MLIRCAPISparseTensor
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
  MODULE_NAME _mlirDialectsTransform
  ADD_TO_PARENT MLIRPythonSources.Dialects.transform
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectTransform.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPIIR
    MLIRCAPITransformDialect
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.IRDL.Pybind
  MODULE_NAME _mlirDialectsIRDL
  ADD_TO_PARENT MLIRPythonSources.Dialects.irdl
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  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}"
  PYTHON_BINDINGS_LIBRARY nanobind
  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}"
    PYTHON_BINDINGS_LIBRARY nanobind
    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}"
  PYTHON_BINDINGS_LIBRARY nanobind
  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}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    LinalgPasses.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPILinalg
)

declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
  MODULE_NAME _mlirDialectsSMT
  ADD_TO_PARENT MLIRPythonSources.Dialects.smt
  ROOT_DIR "${PYTHON_SOURCE_DIR}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    DialectSMT.cpp
    # Headers must be included explicitly so they are installed.
    NanobindUtils.h
  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}"
  PYTHON_BINDINGS_LIBRARY nanobind
  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}"
  PYTHON_BINDINGS_LIBRARY nanobind
  SOURCES
    TransformInterpreter.cpp
  PRIVATE_LINK_LIBS
    LLVMSupport
  EMBED_CAPI_LINK_LIBS
    MLIRCAPITransformDialectTransforms
)

# 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.PythonTestExtensionPybind11
    MODULE_NAME _mlirPythonTestPybind11
    ADD_TO_PARENT MLIRPythonTestSources.Dialects
    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
    PYTHON_BINDINGS_LIBRARY pybind11
    SOURCES
      PythonTestModulePybind11.cpp
    PRIVATE_LINK_LIBS
      LLVMSupport
    EMBED_CAPI_LINK_LIBS
      MLIRCAPIPythonTestDialect
  )
  declare_mlir_python_extension(MLIRPythonTestSources.PythonTestExtensionNanobind
    MODULE_NAME _mlirPythonTestNanobind
    ADD_TO_PARENT MLIRPythonTestSources.Dialects
    ROOT_DIR "${MLIR_SOURCE_DIR}/test/python/lib"
    PYTHON_BINDINGS_LIBRARY nanobind
    SOURCES
      PythonTestModuleNanobind.cpp
    PRIVATE_LINK_LIBS
      LLVMSupport
    EMBED_CAPI_LINK_LIBS
      MLIRCAPIPythonTestDialect
  )
endif()

################################################################################
# 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.
################################################################################

set(MLIRPythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}")
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 ($<BUILD_INTERFACE> $<INSTALL_INTERFACE>)
  # 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()