[MLIR][Python] reland (narrower) type stub generation

.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,89 @@ function(declare_mlir_python_sources name)
  endif()
 endfunction()
 
+# Function: mlir_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: List of paths to add to PYTHONPATH for stubgen.
+# PATTERN_FILE: (Optional) Pattern file (see https://nanobind.readthedocs.io/en/latest/typing.html#pattern-files).
+# Outputs:
+# NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
+function(mlir_generate_type_stubs)
+ cmake_parse_arguments(ARG
+ ""
+ "MODULE_NAME;OUTPUT_DIR;PATTERN_FILE"
+ "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 "mlir_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}")
+ if(ARG_PATTERN_FILE)
+ list(APPEND _nb_stubgen_cmd "-p;${ARG_PATTERN_FILE}")
+ list(APPEND ARG_DEPENDS_TARGETS "${ARG_PATTERN_FILE}")
+ endif()
+
+ 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()
+
+ message(DEBUG "Generating type-stubs outputs ${_generated_type_stubs}")
+ add_custom_command(
+ OUTPUT ${_generated_type_stubs}
+ COMMAND ${_nb_stubgen_cmd}
+ WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
+ DEPENDS "${ARG_DEPENDS_TARGETS}"
+ DEPFILE "${_depfile}"
+ COMMENT "Generating type stubs using: ${_nb_stubgen_cmd}"
+ )
+
+ 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 +761,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/cmake/modules/MLIRDetectPythonEnv.cmake

@@ -54,7 +54,7 @@ macro(mlir_configure_python_dev_packages)
  "extension = '${PYTHON_MODULE_EXTENSION}")
 
  mlir_detect_nanobind_install()
- find_package(nanobind 2.4 CONFIG REQUIRED)
+ find_package(nanobind 2.9 CONFIG REQUIRED)
  message(STATUS "Found nanobind v${nanobind_VERSION}: ${nanobind_INCLUDE_DIR}")
  message(STATUS "Python prefix = '${PYTHON_MODULE_PREFIX}', "
  "suffix = '${PYTHON_MODULE_SUFFIX}', "

mlir/examples/standalone/CMakeLists.txt

@@ -6,9 +6,9 @@ set(CMAKE_BUILD_WITH_INSTALL_NAME_DIR ON)
 set(CMAKE_CXX_STANDARD 17 CACHE STRING "C++ standard to conform to")
 
 if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
+ set(EXTERNAL_PROJECT_BUILD FALSE)
  find_package(MLIR REQUIRED CONFIG)
 
-
  # Define the default argument to use by `lit` when testing.
  set(LLVM_LIT_ARGS "-sv" CACHE STRING "Default options for lit")
 
@@ -27,11 +27,19 @@ if(CMAKE_SOURCE_DIR STREQUAL CMAKE_CURRENT_SOURCE_DIR)
  include(AddMLIR)
  include(HandleLLVMOptions)
 else()
+ set(EXTERNAL_PROJECT_BUILD TRUE)
  # Build via external projects mechanism
+ set(LLVM_RUNTIME_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/bin)
+ set(LLVM_LIBRARY_OUTPUT_INTDIR ${CMAKE_BINARY_DIR}/tools/Standalone/lib)
+
+ set(MLIR_BINARY_DIR ${CMAKE_BINARY_DIR})
  set(MLIR_MAIN_SRC_DIR ${LLVM_MAIN_SRC_DIR}/../mlir)
  set(MLIR_INCLUDE_DIR ${MLIR_MAIN_SRC_DIR}/include)
  set(MLIR_GENERATED_INCLUDE_DIR ${LLVM_BINARY_DIR}/tools/mlir/include)
  set(MLIR_INCLUDE_DIRS "${MLIR_INCLUDE_DIR};${MLIR_GENERATED_INCLUDE_DIR}")
+
+ list(APPEND CMAKE_MODULE_PATH "${MLIR_MAIN_SRC_DIR}/cmake/modules")
+ list(APPEND CMAKE_MODULE_PATH "${LLVM_MAIN_SRC_DIR}/cmake/modules")
 endif()
 
 if(MLIR_ENABLE_BINDINGS_PYTHON)
@@ -54,8 +62,9 @@ if(MLIR_ENABLE_BINDINGS_PYTHON)
  message(STATUS "Enabling Python API")
  include(MLIRDetectPythonEnv)
  mlir_configure_python_dev_packages()
+ # Note: for EXTERNAL_PROJECT_BUILD this must be set from the command line.
  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

@@ -18,6 +18,7 @@ declare_mlir_dialect_python_bindings(
  SOURCES
  dialects/standalone_pybind11.py
  dialects/standalone_nanobind.py
+ _mlir_libs/_standaloneDialectsNanobind/py.typed
  DIALECT_NAME standalone)
 
 
@@ -68,16 +69,85 @@ add_mlir_python_common_capi_library(StandalonePythonCAPI
 # Instantiation of all Python modules
 ################################################################################
 
+set(StandalonePythonModules_ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}")
+
+# Everything here is very tightly coupled. See the ample descriptions at the bottom of
+# mlir/python/CMakeLists.txt.
+
+# For a non-external projects build (e.g., installed distro) the type gen targets for the core _mlir module
+# need to be re-declared. On the contrary, for an external projects build, the MLIRPythonExtension.Core.type_stub_gen
+# target already exists and can just be added to DECLARED_SOURCES (see below).
+if(NOT EXTERNAL_PROJECT_BUILD)
+ set(_core_type_stub_sources
+ _mlir/__init__.pyi
+ _mlir/ir.pyi
+ _mlir/passmanager.pyi
+ _mlir/rewrite.pyi
+ )
+ get_target_property(_core_extension_srcs MLIRPythonExtension.Core INTERFACE_SOURCES)
+ mlir_generate_type_stubs(
+ MODULE_NAME _mlir
+ DEPENDS_TARGETS StandalonePythonModules.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 "${StandalonePythonModules_ROOT_PREFIX}/_mlir_libs"
+ )
+ set(_mlir_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}")
+
+ list(TRANSFORM _core_type_stub_sources PREPEND "_mlir_libs/")
+ declare_mlir_python_sources(
+ StandalonePythonExtension.Core.type_stub_gen
+ ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
+ ADD_TO_PARENT StandalonePythonSources
+ SOURCES "${_core_type_stub_sources}"
+ )
+endif()
+
+get_target_property(_standalone_extension_srcs StandalonePythonSources.NanobindExtension INTERFACE_SOURCES)
+mlir_generate_type_stubs(
+ MODULE_NAME mlir_standalone._mlir_libs._standaloneDialectsNanobind
+ DEPENDS_TARGETS
+ StandalonePythonModules.extension._mlir.dso
+ StandalonePythonModules.extension._standaloneDialectsNanobind.dso
+ OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
+ OUTPUTS
+ _standaloneDialectsNanobind/__init__.pyi
+ _standaloneDialectsNanobind/standalone.pyi
+ DEPENDS_TARGET_SRC_DEPS "${_standalone_extension_srcs}"
+ IMPORT_PATHS "${StandalonePythonModules_ROOT_PREFIX}/.."
+)
+set(_standaloneDialectsNanobind_typestub_gen_target "${NB_STUBGEN_CUSTOM_TARGET}")
+
+declare_mlir_python_sources(
+ StandalonePythonSources.type_stub_gen
+ ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
+ ADD_TO_PARENT StandalonePythonSources
+ SOURCES
+ _mlir_libs/_standaloneDialectsNanobind/__init__.pyi
+ _mlir_libs/_standaloneDialectsNanobind/standalone.pyi
+)
+set(_declared_sources
+ StandalonePythonSources
+ # TODO: Remove this in favor of showing fine grained registration once
+ # available.
+ MLIRPythonExtension.RegisterEverything
+ MLIRPythonSources.Core
+ MLIRPythonSources.Dialects.builtin
+)
+# For an external projects build, the MLIRPythonExtension.Core.type_stub_gen
+# target already exists and can just be added to DECLARED_SOURCES.
+if(EXTERNAL_PROJECT_BUILD)
+ list(APPEND _declared_sources MLIRPythonExtension.Core.type_stub_gen)
+endif()
 add_mlir_python_modules(StandalonePythonModules
- ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
+ ROOT_PREFIX "${StandalonePythonModules_ROOT_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.Core
- MLIRPythonSources.Dialects.builtin
+ DECLARED_SOURCES "${_declared_sources}"
  COMMON_CAPI_LINK_LIBS
  StandalonePythonCAPI
- )
+)
+if(NOT EXTERNAL_PROJECT_BUILD)
+ add_dependencies(StandalonePythonModules "${_mlir_typestub_gen_target}")
+endif()
+add_dependencies(StandalonePythonModules "${_standaloneDialectsNanobind_typestub_gen_target}")

mlir/examples/standalone/python/StandaloneExtensionNanobind.cpp

@@ -30,5 +30,9 @@ NB_MODULE(_standaloneDialectsNanobind, m) {
  mlirDialectHandleLoadDialect(handle, context);
  }
  },
- nb::arg("context").none() = nb::none(), nb::arg("load") = true);
+ nb::arg("context").none() = nb::none(), nb::arg("load") = true,
+ // clang-format off
+ nb::sig("def register_dialect(context: " MAKE_MLIR_PYTHON_QUALNAME("ir.Context") ", load: bool = True) -> None")
+ // clang-format on
+ );
 }

mlir/include/mlir/Bindings/Python/NanobindAdaptors.h

@@ -512,8 +512,13 @@ class mlir_attribute_subclass : public pure_subclass {
  .attr("replace")(superCls.attr("__name__"), captureTypeName);
  });
  if (getTypeIDFunction) {
- def_staticmethod("get_static_typeid",
- [getTypeIDFunction]() { return getTypeIDFunction(); });
+ def_staticmethod(
+ "get_static_typeid",
+ [getTypeIDFunction]() { return getTypeIDFunction(); },
+ // clang-format off
+ nanobind::sig("def get_static_typeid() -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID"))
+ // clang-format on
+ );
  nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
  .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
  getTypeIDFunction())(nanobind::cpp_function(
@@ -582,8 +587,9 @@ class mlir_type_subclass : public pure_subclass {
 
  // 'isinstance' method.
  static const char kIsinstanceSig[] =
- "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME(
- "ir") ".Type) -> bool";
+ // clang-format off
+ "def isinstance(other_type: " MAKE_MLIR_PYTHON_QUALNAME("ir.Type") ") -> bool";
+ // clang-format on
  def_staticmethod(
  "isinstance",
  [isaFunction](MlirType other) { return isaFunction(other); },
@@ -599,8 +605,13 @@ class mlir_type_subclass : public pure_subclass {
  // `def_property_readonly_static` is not available in `pure_subclass` and
  // we do not want to introduce the complexity that pybind uses to
  // implement it.
- def_staticmethod("get_static_typeid",
- [getTypeIDFunction]() { return getTypeIDFunction(); });
+ def_staticmethod(
+ "get_static_typeid",
+ [getTypeIDFunction]() { return getTypeIDFunction(); },
+ // clang-format off
+ nanobind::sig("def get_static_typeid() -> " MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID"))
+ // clang-format on
+ );
  nanobind::module_::import_(MAKE_MLIR_PYTHON_QUALNAME("ir"))
  .attr(MLIR_PYTHON_CAPI_TYPE_CASTER_REGISTER_ATTR)(
  getTypeIDFunction())(nanobind::cpp_function(
@@ -665,8 +676,9 @@ class mlir_value_subclass : public pure_subclass {
 
  // 'isinstance' method.
  static const char kIsinstanceSig[] =
- "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME(
- "ir") ".Value) -> bool";
+ // clang-format off
+ "def isinstance(other_value: " MAKE_MLIR_PYTHON_QUALNAME("ir.Value") ") -> bool";
+ // clang-format on
  def_staticmethod(
  "isinstance",
  [isaFunction](MlirValue other) { return isaFunction(other); },

mlir/lib/Bindings/Python/DialectPDL.cpp

@@ -68,6 +68,8 @@ static void populateDialectPDLSubmodule(const nanobind::module_ &m) {
  rangeType.def_property_readonly(
  "element_type",
  [](MlirType type) { return mlirPDLRangeTypeGetElementType(type); },
+ nb::sig(
+ "def element_type(self) -> " MAKE_MLIR_PYTHON_QUALNAME("ir.Type")),
  "Get the element type.");
 
  //===-------------------------------------------------------------------===//