Revert "[MLIR][Python] fix stubgen (llvm#157583)"

This reverts commit d02c7ae.

Author: rkayaith

mlir/cmake/modules/AddMLIRPython.cmake

@@ -111,11 +111,6 @@ endfunction()
 # 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)
-  cmake_parse_arguments(ARG
-    ""
-    ""
-    "OUTPUTS"
-    ${ARGN})
   if(EXISTS ${nanobind_DIR}/../src/stubgen.py)
     set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py")
   elseif(EXISTS ${nanobind_DIR}/../stubgen.py)
@@ -140,9 +135,9 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP
       --output-dir
       "${OUTPUT_DIR}")
 
-  list(TRANSFORM ARG_OUTPUTS PREPEND "${OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
+  set(NB_STUBGEN_OUTPUT "${OUTPUT_DIR}/${MODULE_NAME}.pyi")
   add_custom_command(
-    OUTPUT ${_generated_type_stubs}
+    OUTPUT ${NB_STUBGEN_OUTPUT}
     COMMAND ${NB_STUBGEN_CMD}
     WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
     DEPENDS
@@ -151,7 +146,7 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP
       "${DEPENDS_TARGET}"
   )
   set(_name "MLIRPythonModuleStubs_${_module}")
-  add_custom_target("${_name}" ALL DEPENDS ${_generated_type_stubs})
+  add_custom_target("${_name}" ALL DEPENDS ${NB_STUBGEN_OUTPUT})
   set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
 endfunction()
 
@@ -171,12 +166,12 @@ 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: Enable type stub generation.
 function(declare_mlir_python_extension name)
   cmake_parse_arguments(ARG
-    ""
+    "GENERATE_TYPE_STUBS"
     "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY"
-    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS;GENERATE_TYPE_STUBS"
+    "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
     ${ARGN})
 
   if(NOT ARG_ROOT_DIR)
@@ -307,25 +302,15 @@ function(add_mlir_python_modules name)
           ${_module_name}
           ${_extension_target}
           ${name}
-          "${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs"
-          OUTPUTS "${_generate_type_stubs}"
+          "${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs/_mlir"
         )
-        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_SOURCE_DIR}/mlir/_mlir_libs"
+          "${MLIR_PYTHON_PACKAGE_PREFIX}.${_module_name}_type_stub_gen"
+          ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
           ADD_TO_PARENT "${sources_target}"
-          SOURCES "${_generate_type_stubs}"
+          SOURCES_GLOB "_mlir_libs/${_module_name}/**/*.pyi"
         )
-        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})
+        add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")
       endif()
     else()
       message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")

mlir/python/CMakeLists.txt

@@ -506,10 +506,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
@@ -532,7 +528,6 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
     MLIRCAPITransforms
     MLIRCAPIRegisterEverything
   GENERATE_TYPE_STUBS
-    "_mlirRegisterEverything.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -548,7 +543,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
     MLIRCAPIIR
     MLIRCAPILinalg
   GENERATE_TYPE_STUBS
-    "_mlirDialectsLinalg.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -564,7 +558,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
     MLIRCAPIIR
     MLIRCAPIGPU
   GENERATE_TYPE_STUBS
-    "_mlirDialectsGPU.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -580,7 +573,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
     MLIRCAPIIR
     MLIRCAPILLVM
   GENERATE_TYPE_STUBS
-    "_mlirDialectsLLVM.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -596,7 +588,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
     MLIRCAPIIR
     MLIRCAPIQuant
   GENERATE_TYPE_STUBS
-    "_mlirDialectsQuant.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -612,7 +603,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
     MLIRCAPIIR
     MLIRCAPINVGPU
   GENERATE_TYPE_STUBS
-    "_mlirDialectsNVGPU.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -628,7 +618,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
     MLIRCAPIIR
     MLIRCAPIPDL
   GENERATE_TYPE_STUBS
-    "_mlirDialectsPDL.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -644,7 +633,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
     MLIRCAPIIR
     MLIRCAPISparseTensor
   GENERATE_TYPE_STUBS
-    "_mlirDialectsSparseTensor.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -660,7 +648,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
     MLIRCAPIIR
     MLIRCAPITransformDialect
   GENERATE_TYPE_STUBS
-    "_mlirDialectsTransform.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -675,7 +662,6 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIAsync
   GENERATE_TYPE_STUBS
-    "_mlirAsyncPasses.pyi"
 )
 
 if(MLIR_ENABLE_EXECUTION_ENGINE)
@@ -691,7 +677,6 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIExecutionEngine
     GENERATE_TYPE_STUBS
-      "_mlirExecutionEngine.pyi"
   )
 endif()
 
@@ -707,7 +692,6 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIGPU
   GENERATE_TYPE_STUBS
-    "_mlirGPUPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -722,7 +706,6 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPILinalg
   GENERATE_TYPE_STUBS
-    "_mlirLinalgPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -741,7 +724,6 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
     MLIRCAPISMT
     MLIRCAPIExportSMTLIB
   GENERATE_TYPE_STUBS
-    "_mlirDialectsSMT.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -756,7 +738,6 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPISparseTensor
   GENERATE_TYPE_STUBS
-    "_mlirSparseTensorPasses.pyi"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -771,7 +752,6 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
   EMBED_CAPI_LINK_LIBS
     MLIRCAPITransformDialectTransforms
   GENERATE_TYPE_STUBS
-    "_mlirTransformInterpreter.pyi"
 )
 
 # TODO: Figure out how to put this in the test tree.
@@ -831,7 +811,6 @@ if(MLIR_INCLUDE_TESTS)
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIPythonTestDialect
     GENERATE_TYPE_STUBS
-      "_mlirPythonTestNanobind.pyi"
   )
 endif()
 

mlir/python/mlir/_mlir_libs/.gitignore

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