Revert stubgen-related patches (#157831)

Despite several hotfixes, things remain broken, in particular:

- installation/distribution (`ninja install / install-distribution`);
- downstream projects with bindings exposed.

See
https://github.com/llvm/llvm-project/pull/157583#issuecomment-3274164357
for more details.

Reverts #155741, #157583, #157697. Let's make sure things are fixed and
re-land as a unit.

mlir/python/CMakeLists.txt

@@ -1,9 +1,5 @@
 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.
 ################################################################################
@@ -27,6 +23,11 @@ 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
@@ -42,6 +43,7 @@ declare_mlir_python_sources(MLIRPythonSources.ExecutionEngine
   ADD_TO_PARENT MLIRPythonSources
   SOURCES
     execution_engine.py
+    _mlir_libs/_mlirExecutionEngine.pyi
   SOURCES_GLOB
     runtime/*.py
 )
@@ -193,6 +195,7 @@ declare_mlir_dialect_python_bindings(
   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"
@@ -364,7 +367,8 @@ declare_mlir_python_sources(
   ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
   GEN_ENUM_BINDINGS
   SOURCES
-    dialects/quant.py)
+    dialects/quant.py
+    _mlir_libs/_mlir/dialects/quant.pyi)
 
 declare_mlir_dialect_python_bindings(
   ADD_TO_PARENT MLIRPythonSources.Dialects
@@ -380,6 +384,7 @@ declare_mlir_dialect_python_bindings(
   TD_FILE dialects/PDLOps.td
   SOURCES
     dialects/pdl.py
+    _mlir_libs/_mlir/dialects/pdl.pyi
   DIALECT_NAME pdl)
 
 declare_mlir_dialect_python_bindings(
@@ -505,11 +510,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
 
     # Dialects
     MLIRCAPIFunc
-  GENERATE_TYPE_STUBS
-    "_mlir/__init__.pyi"
-    "_mlir/ir.pyi"
-    "_mlir/passmanager.pyi"
-    "_mlir/rewrite.pyi"
 )
 
 # This extension exposes an API to register all dialects, extensions, and passes
@@ -531,8 +531,6 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
     MLIRCAPIConversion
     MLIRCAPITransforms
     MLIRCAPIRegisterEverything
-  GENERATE_TYPE_STUBS
-    "_mlirRegisterEverything.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -547,8 +545,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPILinalg
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsLinalg.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -563,8 +559,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPIGPU
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsGPU.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -579,8 +573,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPILLVM
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsLLVM.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -595,8 +587,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPIQuant
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsQuant.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -611,8 +601,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPINVGPU
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsNVGPU.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -627,8 +615,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPIPDL
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsPDL.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -643,8 +629,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPISparseTensor
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsSparseTensor.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -659,8 +643,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIIR
     MLIRCAPITransformDialect
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsTransform.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -674,8 +656,6 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIAsync
-  GENERATE_TYPE_STUBS
-    "_mlirAsyncPasses.pyi"
 )
 
 if(MLIR_ENABLE_EXECUTION_ENGINE)
@@ -690,8 +670,6 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
       LLVMSupport
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIExecutionEngine
-    GENERATE_TYPE_STUBS
-      "_mlirExecutionEngine.pyi"
   )
 endif()
 
@@ -706,8 +684,6 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIGPU
-  GENERATE_TYPE_STUBS
-    "_mlirGPUPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -721,8 +697,6 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPILinalg
-  GENERATE_TYPE_STUBS
-    "_mlirLinalgPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -740,8 +714,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
     MLIRCAPIIR
     MLIRCAPISMT
     MLIRCAPIExportSMTLIB
-  GENERATE_TYPE_STUBS
-    "_mlirDialectsSMT.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -755,8 +727,6 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPISparseTensor
-  GENERATE_TYPE_STUBS
-    "_mlirSparseTensorPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -770,8 +740,6 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
     LLVMSupport
   EMBED_CAPI_LINK_LIBS
     MLIRCAPITransformDialectTransforms
-  GENERATE_TYPE_STUBS
-    "_mlirTransformInterpreter.pyi"
 )
 
 # TODO: Figure out how to put this in the test tree.
@@ -830,8 +798,6 @@ if(MLIR_INCLUDE_TESTS)
       LLVMSupport
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIPythonTestDialect
-    GENERATE_TYPE_STUBS
-      "_mlirPythonTestNanobind.pyi"
   )
 endif()
 
@@ -851,7 +817,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}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/_mlir_libs"
+  OUTPUT_DIRECTORY "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir/_mlir_libs"
   RELATIVE_INSTALL_ROOT "../../../.."
   DECLARED_HEADERS
     MLIRPythonCAPI.HeaderSources
@@ -880,7 +846,7 @@ endif()
 ################################################################################
 
 add_mlir_python_modules(MLIRPythonModules
-  ROOT_PREFIX "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
+  ROOT_PREFIX "${MLIR_BINARY_DIR}/python_packages/mlir_core/mlir"
   INSTALL_PREFIX "${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}"
   DECLARED_SOURCES
     MLIRPythonSources

mlir/python/mlir/_mlir_libs/.gitignore

@@ -1,2 +0,0 @@
-_mlir/**/*.pyi
-*.pyi

mlir/python/mlir/_mlir_libs/_mlir/__init__.pyi

@@ -0,0 +1,12 @@
+
+globals: "_Globals"
+
+class _Globals:
+    dialect_search_modules: list[str]
+    def _register_dialect_impl(self, dialect_namespace: str, dialect_class: type) -> None: ...
+    def _register_operation_impl(self, operation_name: str, operation_class: type) -> None: ...
+    def append_dialect_search_prefix(self, module_name: str) -> None: ...
+    def _check_dialect_module_loaded(self, dialect_namespace: str) -> bool: ...
+
+def register_dialect(dialect_class: type) -> type: ...
+def register_operation(dialect_class: type, *, replace: bool = ...) -> type: ...

mlir/python/mlir/_mlir_libs/_mlir/dialects/pdl.pyi

@@ -0,0 +1,63 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+
+from mlir.ir import Type, Context
+
+__all__ = [
+    'PDLType',
+    'AttributeType',
+    'OperationType',
+    'RangeType',
+    'TypeType',
+    'ValueType',
+]
+
+
+class PDLType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+
+class AttributeType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Context | None = None) -> AttributeType: ...
+
+
+class OperationType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Context | None = None) -> OperationType: ...
+
+
+class RangeType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(element_type: Type) -> RangeType: ...
+
+  @property
+  def element_type(self) -> Type: ...
+
+
+class TypeType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Context | None = None) -> TypeType: ...
+
+
+class ValueType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Context | None = None) -> ValueType: ...

mlir/python/mlir/_mlir_libs/_mlir/dialects/quant.pyi

@@ -0,0 +1,142 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+
+from mlir.ir import DenseElementsAttr, Type
+
+__all__ = [
+  "QuantizedType",
+  "AnyQuantizedType",
+  "UniformQuantizedType",
+  "UniformQuantizedPerAxisType",
+  "CalibratedQuantizedType",
+]
+
+class QuantizedType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def default_minimum_for_integer(is_signed: bool, integral_width: int) -> int:
+    ...
+
+  @staticmethod
+  def default_maximum_for_integer(is_signed: bool, integral_width: int) -> int:
+    ...
+
+  @property
+  def expressed_type(self) -> Type: ...
+
+  @property
+  def flags(self) -> int: ...
+
+  @property
+  def is_signed(self) -> bool: ...
+
+  @property
+  def storage_type(self) -> Type: ...
+
+  @property
+  def storage_type_min(self) -> int: ...
+
+  @property
+  def storage_type_max(self) -> int: ...
+
+  @property
+  def storage_type_integral_width(self) -> int: ...
+
+  def is_compatible_expressed_type(self, candidate: Type) -> bool: ...
+
+  @property
+  def quantized_element_type(self) -> Type: ...
+
+  def cast_from_storage_type(self, candidate: Type) -> Type: ...
+
+  @staticmethod
+  def cast_to_storage_type(type: Type) -> Type: ...
+
+  def cast_from_expressed_type(self, candidate: Type) -> Type: ...
+
+  @staticmethod
+  def cast_to_expressed_type(type: Type) -> Type: ...
+
+  def cast_expressed_to_storage_type(self, candidate: Type) -> Type: ...
+
+
+class AnyQuantizedType(QuantizedType):
+
+  @classmethod
+  def get(cls, flags: int, storage_type: Type, expressed_type: Type,
+          storage_type_min: int, storage_type_max: int) -> Type:
+    ...
+
+
+class UniformQuantizedType(QuantizedType):
+
+  @classmethod
+  def get(cls, flags: int, storage_type: Type, expressed_type: Type,
+          scale: float, zero_point: int, storage_type_min: int,
+          storage_type_max: int) -> Type: ...
+
+  @property
+  def scale(self) -> float: ...
+
+  @property
+  def zero_point(self) -> int: ...
+
+  @property
+  def is_fixed_point(self) -> bool: ...
+
+
+class UniformQuantizedPerAxisType(QuantizedType):
+
+  @classmethod
+  def get(cls, flags: int, storage_type: Type, expressed_type: Type,
+          scales: list[float], zero_points: list[int], quantized_dimension: int,
+          storage_type_min: int, storage_type_max: int):
+    ...
+
+  @property
+  def scales(self) -> list[float]: ...
+
+  @property
+  def zero_points(self) -> list[int]: ...
+
+  @property
+  def quantized_dimension(self) -> int: ...
+
+  @property
+  def is_fixed_point(self) -> bool: ...
+
+class UniformQuantizedSubChannelType(QuantizedType):
+
+  @classmethod
+  def get(cls, flags: int, storage_type: Type, expressed_type: Type,
+          scales: DenseElementsAttr, zero_points: DenseElementsAttr,
+          quantized_dimensions: list[int], block_sizes: list[int],
+          storage_type_min: int, storage_type_max: int):
+    ...
+
+  @property
+  def quantized_dimensions(self) -> list[int]: ...
+
+  @property
+  def block_sizes(self) -> list[int]: ...
+
+  @property
+  def scales(self) -> DenseElementsAttr: ...
+
+  @property
+  def zero_points(self) -> DenseElementsAttr: ...
+
+def CalibratedQuantizedType(QuantizedType):
+
+  @classmethod
+  def get(cls, expressed_type: Type, min: float, max: float): ...
+
+  @property
+  def min(self) -> float: ...
+
+  @property
+  def max(self) -> float: ...

mlir/python/mlir/_mlir_libs/_mlir/dialects/transform/__init__.pyi

@@ -0,0 +1,25 @@
+#  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+#  See https://llvm.org/LICENSE.txt for license information.
+#  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+
+from mlir.ir import Type, Context
+
+
+class AnyOpType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(context: Context | None = None) -> AnyOpType: ...
+
+
+class OperationType(Type):
+  @staticmethod
+  def isinstance(type: Type) -> bool: ...
+
+  @staticmethod
+  def get(operation_name: str, context: Context | None = None) -> OperationType: ...
+
+  @property
+  def operation_name(self) -> str: ...

mlir/python/mlir/_mlir_libs/_mlir/passmanager.pyi

@@ -0,0 +1,36 @@
+# Originally imported via:
+#   stubgen {...} -m mlir._mlir_libs._mlir.passmanager
+# Local modifications:
+#   * Relative imports for cross-module references.
+#   * Add __all__
+
+
+from . import ir as _ir
+
+__all__ = [
+    "PassManager",
+]
+
+class PassManager:
+    def __init__(self, context: _ir.Context | None = None) -> None: ...
+    def _CAPICreate(self) -> object: ...
+    def _testing_release(self) -> None: ...
+    def enable_ir_printing(
+        self,
+        print_before_all: bool = False,
+        print_after_all: bool = True,
+        print_module_scope: bool = False,
+        print_after_change: bool = False,
+        print_after_failure: bool = False,
+        large_elements_limit: int | None = None,
+        large_resource_limit: int | None = None,
+        enable_debug_info: bool = False,
+        print_generic_op_form: bool = False,
+        tree_printing_dir_path: str | None = None,
+    ) -> None: ...
+    def enable_verifier(self, enable: bool) -> None: ...
+    @staticmethod
+    def parse(pipeline: str, context: _ir.Context | None = None) -> PassManager: ...
+    def run(self, module: _ir._OperationBase) -> None: ...
+    @property
+    def _CAPIPtr(self) -> object: ...

mlir/python/mlir/_mlir_libs/_mlirExecutionEngine.pyi

@@ -0,0 +1,24 @@
+# Originally imported via:
+#   stubgen {...} -m mlir._mlir_libs._mlirExecutionEngine
+# Local modifications:
+#   * Relative imports for cross-module references.
+#   * Add __all__
+
+from collections.abc import Sequence
+
+from ._mlir import ir as _ir
+
+__all__ = [
+    "ExecutionEngine",
+]
+
+class ExecutionEngine:
+    def __init__(self, module: _ir.Module, opt_level: int = 2, shared_libs: Sequence[str] = ...) -> None: ...
+    def _CAPICreate(self) -> object: ...
+    def _testing_release(self) -> None: ...
+    def dump_to_object_file(self, file_name: str) -> None: ...
+    def raw_lookup(self, func_name: str) -> int: ...
+    def raw_register_runtime(self, name: str, callback: object) -> None: ...
+    def init() -> None: ...
+    @property
+    def _CAPIPtr(self) -> object: ...

mlir/python/requirements.txt

@@ -1,7 +1,6 @@
-nanobind>=2.9, <3.0
+nanobind>=2.4, <3.0
 numpy>=1.19.5, <=2.1.2
 pybind11>=2.10.0, <=2.13.6
 PyYAML>=5.4.0, <=6.0.1
 ml_dtypes>=0.1.0, <=0.6.0; python_version<"3.13"   # provides several NumPy dtype extensions, including the bf16
-ml_dtypes>=0.5.0, <=0.6.0; python_version>="3.13"
-typing_extensions>=4.12.2
+ml_dtypes>=0.5.0, <=0.6.0; python_version>="3.13"