Update AddMLIRPython.cmake

Author: makslevental

mlir/cmake/modules/AddMLIRPython.cmake

@@ -99,7 +99,7 @@ function(declare_mlir_python_sources name)
   endif()
 endfunction()
 
-function(generate_type_stubs module_name depends_target output_dir)
+function(generate_type_stubs module_name depends_target mlir_depends_target output_dir)
   if(EXISTS ${nanobind_DIR}/../src/stubgen.py)
     set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py")
   elseif(EXISTS ${nanobind_DIR}/../stubgen.py)
@@ -108,11 +108,12 @@ function(generate_type_stubs module_name depends_target output_dir)
       message(FATAL_ERROR "generate_type_stubs(): could not locate 'stubgen.py'!")
   endif()
 
+  set(_module "${MLIR_PYTHON_PACKAGE_PREFIX}._mlir_libs.${module_name}")
   set(NB_STUBGEN_CMD
       "${Python_EXECUTABLE}"
       "${NB_STUBGEN}"
       --module
-      "${MLIR_PYTHON_PACKAGE_PREFIX}._mlir_libs.${module_name}"
+      "${_module}"
       -i
       "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.."
       --recursive
@@ -125,8 +126,8 @@ function(generate_type_stubs module_name depends_target output_dir)
     OUTPUT ${NB_STUBGEN_OUTPUT}
     COMMAND ${NB_STUBGEN_CMD}
     WORKING_DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}"
-    DEPENDS ${depends_target})
-  set(_name "MLIRPythonModuleStubs_${module_name}")
+    DEPENDS "${mlir_depends_target}" "${depends_target}")
+  set(_name "MLIRPythonModuleStubs_${_module}")
   add_custom_target("${_name}" ALL DEPENDS ${NB_STUBGEN_OUTPUT})
   set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
 endfunction()
@@ -278,9 +279,11 @@ function(add_mlir_python_modules name)
       generate_type_stubs(
         ${_module_name}
         ${_extension_target}
+        "${modules_target}.extension._mlir.dso"
         "${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs/_mlir"
       )
-      declare_mlir_python_sources("_${_module_name}_type_stub_gen"
+      declare_mlir_python_sources(
+        "${MLIR_PYTHON_PACKAGE_PREFIX}.${_module_name}_type_stub_gen"
         ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
         ADD_TO_PARENT "${sources_target}"
         SOURCES_GLOB "_mlir_libs/${_module_name}/**/*.pyi"

mlir/examples/standalone/CMakeLists.txt

@@ -52,6 +52,8 @@ add_subdirectory(include)
 add_subdirectory(lib)
 if(MLIR_ENABLE_BINDINGS_PYTHON)
   message(STATUS "Enabling Python API")
+  set(MLIR_PYTHON_PACKAGE_PREFIX "mlir_standalone" CACHE STRING "" FORCE)
+  set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/standalone/mlir_standalone" CACHE STRING "" FORCE)
   add_subdirectory(python)
 endif()
 add_subdirectory(test)

mlir/examples/standalone/python/CMakeLists.txt

@@ -48,8 +48,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
 
 add_mlir_python_common_capi_library(StandalonePythonCAPI
   INSTALL_COMPONENT StandalonePythonModules
-  INSTALL_DESTINATION python_packages/standalone/mlir_standalone/_mlir_libs
-  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/standalone/mlir_standalone/_mlir_libs"
+  INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
+  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_SOURCES
     StandalonePythonSources
@@ -64,14 +64,14 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
 ################################################################################
 
 add_mlir_python_modules(StandalonePythonModules
-  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/standalone/mlir_standalone"
-  INSTALL_PREFIX "python_packages/standalone/mlir_standalone"
+  ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
+  INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
   DECLARED_SOURCES
     StandalonePythonSources
     # TODO: Remove this in favor of showing fine grained registration once
     # available.
     MLIRPythonExtension.RegisterEverything
-    MLIRPythonSources
+    MLIRPythonSources.Core
   COMMON_CAPI_LINK_LIBS
     StandalonePythonCAPI
   )