[MLIR][Python] fix stubgen for downstream users

mlir/cmake/modules/AddMLIRPython.cmake

@@ -102,57 +102,72 @@ endfunction()
 # Function: generate_type_stubs
 # Turns on automatic type stub generation (via nanobind's stubgen) for extension modules.
 # Arguments:
-#   MODULE_NAME: The name of the extension module as specified in declare_mlir_python_extension.
+#   FQ_MODULE_NAME: The fully-qualified name of the extension module (used for importing in python).
 #   DEPENDS_TARGET: The dso target corresponding to the extension module
 #     (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso)
-#   MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module
+#   CORE_MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module
 #     (e.g., something like StandalonePythonModules.extension._mlir.dso)
 #   OUTPUT_DIR: The root output directory to emit the type stubs into.
+#   OUTPUTS: List of expected outputs.
+#   DEPENDS_TARGET_SRC_DEPS: List of cpp sources for extension library (for generating a DEPFILE).
 # Outputs:
 #   NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
-function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR)
+function(generate_type_stubs )
   cmake_parse_arguments(ARG
     ""
-    ""
-    "OUTPUTS"
+    "FQ_MODULE_NAME;DEPENDS_TARGET;CORE_MLIR_DEPENDS_TARGET;OUTPUT_DIR"
+    "OUTPUTS;DEPENDS_TARGET_SRC_DEPS"
     ${ARGN})
+  # for people doing find_package(nanobind)
   if(EXISTS ${nanobind_DIR}/../src/stubgen.py)
     set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py")
   elseif(EXISTS ${nanobind_DIR}/../stubgen.py)
     set(NB_STUBGEN "${nanobind_DIR}/../stubgen.py")
+  # for people using FetchContent_Declare and FetchContent_MakeAvailable
+  elseif(EXISTS ${nanobind_SOURCE_DIR}/src/stubgen.py)
+    set(NB_STUBGEN "${nanobind_SOURCE_DIR}/src/stubgen.py")
+  elseif(EXISTS ${nanobind_SOURCE_DIR}/stubgen.py)
+    set(NB_STUBGEN "${nanobind_SOURCE_DIR}/stubgen.py")
   else()
     message(FATAL_ERROR "generate_type_stubs(): could not locate 'stubgen.py'!")
   endif()
   file(REAL_PATH "${NB_STUBGEN}" NB_STUBGEN)
 
-  set(_module "${MLIR_PYTHON_PACKAGE_PREFIX}._mlir_libs.${MODULE_NAME}")
   file(REAL_PATH "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.." _import_path)
 
-  set(NB_STUBGEN_CMD
+  set(_nb_stubgen_cmd
       "${Python_EXECUTABLE}"
       "${NB_STUBGEN}"
       --module
-      "${_module}"
+      "${ARG_FQ_MODULE_NAME}"
       -i
       "${_import_path}"
       --recursive
       --include-private
       --output-dir
-      "${OUTPUT_DIR}"
+      "${ARG_OUTPUT_DIR}"
       --quiet)
 
-  list(TRANSFORM ARG_OUTPUTS PREPEND "${OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
+  list(TRANSFORM ARG_OUTPUTS PREPEND "${ARG_OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
+  set(_depfile "${ARG_OUTPUT_DIR}/${ARG_FQ_MODULE_NAME}.d")
+  if ((NOT EXISTS ${_depfile}) AND ARG_DEPENDS_TARGET_SRC_DEPS)
+    list(JOIN ARG_DEPENDS_TARGET_SRC_DEPS " " _depfiles)
+    list(TRANSFORM _generated_type_stubs APPEND ": ${_depfiles}" OUTPUT_VARIABLE _depfiles)
+    list(JOIN _depfiles "\n" _depfiles)
+    file(GENERATE OUTPUT "${_depfile}" CONTENT "${_depfiles}")
+  endif()
   add_custom_command(
     OUTPUT ${_generated_type_stubs}
-    COMMAND ${NB_STUBGEN_CMD}
+    COMMAND ${_nb_stubgen_cmd}
     WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
     DEPENDS
-      "${MLIR_DEPENDS_TARGET}.extension._mlir.dso"
-      "${MLIR_DEPENDS_TARGET}.sources.MLIRPythonSources.Core.Python"
-      "${DEPENDS_TARGET}"
+      "${ARG_CORE_MLIR_DEPENDS_TARGET}.extension._mlir.dso"
+      "${ARG_CORE_MLIR_DEPENDS_TARGET}.sources.MLIRPythonSources.Core.Python"
+      "${ARG_DEPENDS_TARGET}"
+    DEPFILE "${_depfile}"
   )
-  set(_name "MLIRPythonModuleStubs_${_module}")
-  add_custom_target("${_name}" ALL DEPENDS ${_generated_type_stubs})
+  set(_name "${ARG_FQ_MODULE_NAME}.type_stubs")
+  add_custom_target("${_name}" DEPENDS ${_generated_type_stubs})
   set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
 endfunction()
 
@@ -172,7 +187,10 @@ endfunction()
 #     on. These will be collected for all extensions and put into an
 #     aggregate dylib that is linked against.
 #   PYTHON_BINDINGS_LIBRARY: Either pybind11 or nanobind.
-#   GENERATE_TYPE_STUBS: List of generated type stubs expected from stubgen relative to _mlir_libs.
+#   GENERATE_TYPE_STUBS: Either
+#     1. OFF (default)
+#     2. ON if ${MODULE_NAME}.pyi is the only stub
+#     3. A list of generated type stubs expected from stubgen relative to _mlir_libs.
 function(declare_mlir_python_extension name)
   cmake_parse_arguments(ARG
     ""
@@ -190,6 +208,23 @@ function(declare_mlir_python_extension name)
   endif()
 
   add_library(${name} INTERFACE)
+
+  if(NOT ARG_GENERATE_TYPE_STUBS)
+    set(ARG_GENERATE_TYPE_STUBS OFF)
+  endif()
+  if("${ARG_GENERATE_TYPE_STUBS}" STREQUAL "ON")
+    set(ARG_GENERATE_TYPE_STUBS "${ARG_MODULE_NAME}.pyi")
+  endif()
+  if(ARG_GENERATE_TYPE_STUBS)
+    list(TRANSFORM ARG_GENERATE_TYPE_STUBS PREPEND "_mlir_libs/")
+    declare_mlir_python_sources(
+      "${name}.type_stub_gen"
+      ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
+      ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
+      SOURCES "${ARG_GENERATE_TYPE_STUBS}"
+    )
+  endif()
+
   set_target_properties(${name} PROPERTIES
     # Yes: Leading-lowercase property names are load bearing and the recommended
     # way to do this: https://gitlab.kitware.com/cmake/cmake/-/issues/19261
@@ -265,10 +300,12 @@ endfunction()
 #     DAG of source modules is included.
 #   COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every
 #     extension depend on (see mlir_python_add_common_capi_library).
+#   PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX not including trailing `.`.
+#     This is used to determine type stub generation python module names.
 function(add_mlir_python_modules name)
   cmake_parse_arguments(ARG
     ""
-    "ROOT_PREFIX;INSTALL_PREFIX"
+    "ROOT_PREFIX;INSTALL_PREFIX;PACKAGE_PREFIX"
     "COMMON_CAPI_LINK_LIBS;DECLARED_SOURCES"
     ${ARGN})
   # Helper to process an individual target.
@@ -302,31 +339,40 @@ function(add_mlir_python_modules name)
       )
       add_dependencies(${modules_target} ${_extension_target})
       mlir_python_setup_extension_rpath(${_extension_target})
+      # NOTE: `sources_target` (naturally) lists all the sources (it's the INTERFACE
+      # target defined above in declare_mlir_python_extension). It's also the name of the
+      # target that gets exported (i.e., is populated as an INTERFACE IMPORTED library in MLIRTargets.cmake).
+      # This is why all metadata is queried from `sources_target`. On the other hand
+      # `_extension_target` is the actual dylib target that's built just above with `add_mlir_python_extension`.
+      # That's why dependencies are in terms of `_extension_target`.
       get_target_property(_generate_type_stubs ${sources_target} mlir_python_GENERATE_TYPE_STUBS)
       if(_generate_type_stubs)
+        if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL ""))
+          message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX for ${name}")
+        endif()
+        # TL;DR: all paths here are load bearing and annoyingly coupled. Changing paths here
+        # (or related code in declare_mlir_python_extension) will break either in-tree or out-of-tree generation.
+        #
+        # We remove _mlir_libs here because OUTPUT_DIR already includes it.
+        # Specifically OUTPUT_DIR already includes it because that's the actual directory
+        # where we want stubgen to dump the emitted sources. This is load bearing because up above
+        # (in declare_mlir_python_extension) we prefixed all the paths with _mlir_libs and
+        # we specified declare_mlir_python_sources with ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs".
+        #
+        # NOTE: INTERFACE_SOURCES is a genex in the build dir ($<BUILD_INTERFACE> and $<INSTALL_INTERFACE>)
+        # which will be evaluated by file(GENERATE ...). In the install dir it's a conventional path
+        # (see install/lib/cmake/mlir/MLIRTargets.cmake).
+        get_target_property(_extension_srcs ${sources_target} INTERFACE_SOURCES)
+        list(TRANSFORM _generate_type_stubs REPLACE "_mlir_libs/" "")
         generate_type_stubs(
-          ${_module_name}
-          ${_extension_target}
-          ${name}
-          "${CMAKE_CURRENT_BINARY_DIR}/_mlir_libs"
+          FQ_MODULE_NAME "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}"
+          DEPENDS_TARGET ${_extension_target}
+          CORE_MLIR_DEPENDS_TARGET ${name}
+          OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
           OUTPUTS "${_generate_type_stubs}"
+          DEPENDS_TARGET_SRC_DEPS "${_extension_srcs}"
         )
         add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")
-        set(_stubgen_target "${MLIR_PYTHON_PACKAGE_PREFIX}.${_module_name}_type_stub_gen")
-        declare_mlir_python_sources(
-          ${_stubgen_target}
-          ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/_mlir_libs"
-          ADD_TO_PARENT "${sources_target}"
-          SOURCES "${_generate_type_stubs}"
-        )
-        set(_pure_sources_target "${modules_target}.sources.${sources_target}_type_stub_gen")
-        add_mlir_python_sources_target(${_pure_sources_target}
-          INSTALL_COMPONENT ${modules_target}
-          INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
-          OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
-          SOURCES_TARGETS ${_stubgen_target}
-        )
-        add_dependencies(${modules_target} ${_pure_sources_target})
       endif()
     else()
       message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")

mlir/examples/standalone/CMakeLists.txt

@@ -55,7 +55,7 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
   include(MLIRDetectPythonEnv)
   mlir_configure_python_dev_packages()
   set(MLIR_PYTHON_PACKAGE_PREFIX "mlir_standalone" CACHE STRING "" FORCE)
-  set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/standalone/mlir_standalone" CACHE STRING "" FORCE)
+  set(MLIR_BINDINGS_PYTHON_INSTALL_PREFIX "python_packages/standalone/${MLIR_PYTHON_PACKAGE_PREFIX}" CACHE STRING "" FORCE)
   add_subdirectory(python)
 endif()
 add_subdirectory(test)

mlir/examples/standalone/python/CMakeLists.txt

@@ -40,6 +40,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY nanobind
   GENERATE_TYPE_STUBS
+    "_standaloneDialectsNanobind/__init__.pyi"
+    "_standaloneDialectsNanobind/standalone.pyi"
 )
 
 
@@ -77,4 +79,5 @@ add_mlir_python_modules(StandalonePythonModules
     MLIRPythonSources.Dialects.builtin
   COMMON_CAPI_LINK_LIBS
     StandalonePythonCAPI
-  )
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
+)