[MLIR][Python] reland stubgen v2

Author: makslevental

.ci/all_requirements.txt

@@ -194,9 +194,9 @@ ml-dtypes==0.5.1 ; python_version < "3.13" \
     --hash=sha256:d13755f8e8445b3870114e5b6240facaa7cb0c3361e54beba3e07fa912a6e12b \
     --hash=sha256:fd918d4e6a4e0c110e2e05be7a7814d10dc1b95872accbf6512b80a109b71ae1
     # via -r mlir/python/requirements.txt
-nanobind==2.7.0 \
-    --hash=sha256:73b12d0e751d140d6c1bf4b215e18818a8debfdb374f08dc3776ad208d808e74 \
-    --hash=sha256:f9f1b160580c50dcf37b6495a0fd5ec61dc0d95dae5f8004f87dd9ad7eb46b34
+nanobind==2.9.2 \
+    --hash=sha256:c37957ffd5eac7eda349cff3622ecd32e5ee1244ecc912c99b5bc8188bafd16e \
+    --hash=sha256:e7608472de99d375759814cab3e2c94aba3f9ec80e62cfef8ced495ca5c27d6e
     # via -r mlir/python/requirements.txt
 numpy==2.0.2 \
     --hash=sha256:0123ffdaa88fa4ab64835dcbde75dcdf89c453c922f18dced6e27c90d1d0ec5a \
@@ -383,6 +383,10 @@ swig==4.3.1 \
     --hash=sha256:efec16327029f682f649a26da726bb0305be8800bd0f1fa3e81bf0769cf5b476 \
     --hash=sha256:fc496c0d600cf1bb2d91e28d3d6eae9c4301e5ea7a0dec5a4281b5efed4245a8
     # via -r lldb/test/requirements.txt
+typing-extensions==4.15.0 \
+    --hash=sha256:0cea48d173cc12fa28ecabc3b837ea3cf6f38c6d1136f85cbaaf598984861466 \
+    --hash=sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548
+    # via -r mlir/python/requirements.txt
 urllib3==2.5.0 \
     --hash=sha256:3fc47733c7e419d4bc3f6b3dc2b4f890bb743906a30d56ba4a5bfa4bbff92760 \
     --hash=sha256:e6b01673c0fa6a13e374b50871808eb3bf7046c4b125b216f6bf1cc604cff0dc

mlir/cmake/modules/AddMLIRPython.cmake

@@ -99,6 +99,82 @@ function(declare_mlir_python_sources name)
   endif()
 endfunction()
 
+# Function: generate_type_stubs
+# Turns on automatic type stub generation for extension modules.
+# Specifically, performs add_custom_command to run nanobind's stubgen on an extension module.
+#
+# Arguments:
+#   MODULE_NAME: The fully-qualified name of the extension module (used for importing in python).
+#   DEPENDS_TARGETS: List of targets these type stubs depend on being built; usually corresponding to the
+#     specific extension module (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso)
+#     and the core bindings 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).
+#   IMPORT_PATHS:
+# Outputs:
+#   NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
+function(generate_type_stubs)
+  cmake_parse_arguments(ARG
+    ""
+    "MODULE_NAME;OUTPUT_DIR"
+    "IMPORT_PATHS;DEPENDS_TARGETS;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(_import_paths)
+  foreach(_import_path IN LISTS ARG_IMPORT_PATHS)
+    file(REAL_PATH "${_import_path}" _import_path)
+    list(APPEND _import_paths "-i;${_import_path}")
+  endforeach()
+  set(_nb_stubgen_cmd
+      "${Python_EXECUTABLE}"
+      "${NB_STUBGEN}"
+      --module
+      "${ARG_MODULE_NAME}"
+      "${_import_paths}"
+      --recursive
+      --include-private
+      --output-dir
+      "${ARG_OUTPUT_DIR}")
+
+  list(TRANSFORM ARG_OUTPUTS PREPEND "${ARG_OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
+  set(_depfile "${ARG_OUTPUT_DIR}/${ARG_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}
+    WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
+    DEPENDS "${ARG_DEPENDS_TARGETS}"
+    DEPFILE "${_depfile}"
+  )
+
+  list(JOIN ARG_DEPENDS_TARGETS "." _name)
+  set(_name "${_name}.${ARG_MODULE_NAME}.type_stubs")
+  add_custom_target("${_name}" DEPENDS ${_generated_type_stubs})
+  set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
+endfunction()
+
 # Function: declare_mlir_python_extension
 # Declares a buildable python extension from C++ source files. The built
 # module is considered a python source file and included as everything else.
@@ -678,26 +754,28 @@ function(add_mlir_python_extension libname extname)
       # the super project handle compile options as it wishes.
       get_property(NB_LIBRARY_TARGET_NAME TARGET ${libname} PROPERTY LINK_LIBRARIES)
       target_compile_options(${NB_LIBRARY_TARGET_NAME}
-	PRIVATE
-	  -Wall -Wextra -Wpedantic
-	  -Wno-c++98-compat-extra-semi
-	  -Wno-cast-qual
-	  -Wno-covered-switch-default
-	  -Wno-nested-anon-types
-	  -Wno-unused-parameter
-	  -Wno-zero-length-array
-	  ${eh_rtti_enable})
+        PRIVATE
+          -Wall -Wextra -Wpedantic
+          -Wno-c++98-compat-extra-semi
+          -Wno-cast-qual
+          -Wno-covered-switch-default
+          -Wno-deprecated-literal-operator
+          -Wno-nested-anon-types
+          -Wno-unused-parameter
+          -Wno-zero-length-array
+          ${eh_rtti_enable})
 
       target_compile_options(${libname}
-	PRIVATE
-	  -Wall -Wextra -Wpedantic
-	  -Wno-c++98-compat-extra-semi
-	  -Wno-cast-qual
-	  -Wno-covered-switch-default
-	  -Wno-nested-anon-types
-	  -Wno-unused-parameter
-	  -Wno-zero-length-array
-	  ${eh_rtti_enable})
+        PRIVATE
+          -Wall -Wextra -Wpedantic
+          -Wno-c++98-compat-extra-semi
+          -Wno-cast-qual
+          -Wno-covered-switch-default
+          -Wno-deprecated-literal-operator
+          -Wno-nested-anon-types
+          -Wno-unused-parameter
+          -Wno-zero-length-array
+          ${eh_rtti_enable})
     endif()
 
     if(APPLE)

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

@@ -76,4 +76,4 @@ add_mlir_python_modules(StandalonePythonModules
     MLIRPythonSources.Dialects.builtin
   COMMON_CAPI_LINK_LIBS
     StandalonePythonCAPI
-  )
+)

mlir/python/CMakeLists.txt

@@ -1,5 +1,9 @@
 include(AddMLIRPython)
 
+# Specifies that all MLIR packages are co-located under the `mlir_standalone`
+# 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.
 ################################################################################
@@ -23,11 +27,6 @@ declare_mlir_python_sources(MLIRPythonSources.Core.Python
     passmanager.py
     rewrite.py
     dialects/_ods_common.py
-
-    # The main _mlir module has submodules: include stubs from each.
-    _mlir_libs/_mlir/__init__.pyi
-    _mlir_libs/_mlir/ir.pyi
-    _mlir_libs/_mlir/passmanager.pyi
 )
 
 declare_mlir_python_sources(MLIRPythonSources.Core.Python.Extras
@@ -479,28 +478,30 @@ declare_mlir_dialect_python_bindings(
 ################################################################################
 
 set(PYTHON_SOURCE_DIR "${MLIR_SOURCE_DIR}/lib/Bindings/Python")
+set(MLIRPythonExtension.Core.CPPSources
+  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
+)
 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
+  SOURCES "${MLIRPythonExtension.Core.CPPSources}"
   PRIVATE_LINK_LIBS
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
@@ -511,6 +512,18 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     # Dialects
     MLIRCAPIFunc
 )
+set(_type_stub_sources
+  _mlir_libs/_mlir/__init__.pyi
+  _mlir_libs/_mlir/ir.pyi
+  _mlir_libs/_mlir/passmanager.pyi
+  _mlir_libs/_mlir/rewrite.pyi
+)
+declare_mlir_python_sources(
+  MLIRPythonExtension.Core.type_stub_gen
+  ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
+  ADD_TO_PARENT MLIRPythonSources.Core
+  SOURCES "${_type_stub_sources}"
+)
 
 # 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
@@ -663,7 +676,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
     MODULE_NAME _mlirExecutionEngine
     ADD_TO_PARENT MLIRPythonSources.ExecutionEngine
     ROOT_DIR "${PYTHON_SOURCE_DIR}"
-  PYTHON_BINDINGS_LIBRARY nanobind
+    PYTHON_BINDINGS_LIBRARY nanobind
     SOURCES
       ExecutionEngineModule.cpp
     PRIVATE_LINK_LIBS
@@ -817,7 +830,7 @@ endif()
 add_mlir_python_common_capi_library(MLIRPythonCAPI
   INSTALL_COMPONENT MLIRPythonModules
   INSTALL_DESTINATION "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
-  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
+  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_HEADERS
     MLIRPythonCAPI.HeaderSources
@@ -845,8 +858,9 @@ endif()
 # This must come last.
 ################################################################################
 
+set(MLIRPythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}")
 add_mlir_python_modules(MLIRPythonModules
-  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
+  ROOT_PREFIX ${MLIRPythonModules_ROOT_PREFIX}
   INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
   DECLARED_SOURCES
     MLIRPythonSources
@@ -855,3 +869,14 @@ add_mlir_python_modules(MLIRPythonModules
   COMMON_CAPI_LINK_LIBS
     MLIRPythonCAPI
 )
+list(TRANSFORM _type_stub_sources REPLACE "_mlir_libs/" "")
+generate_type_stubs(
+  MODULE_NAME _mlir_libs._mlir
+  DEPENDS_TARGETS
+    MLIRPythonModules.extension._mlir.dso
+  OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
+  OUTPUTS "${_type_stub_sources}"
+  DEPENDS_TARGET_SRC_DEPS "${MLIRPythonExtension.Core.CPPSources}"
+  IMPORT_PATHS "${MLIRPythonModules_ROOT_PREFIX}"
+)
+add_dependencies(MLIRPythonModules "${NB_STUBGEN_CUSTOM_TARGET}")