god what a pain

Author: makslevental

mlir/cmake/modules/AddMLIRPython.cmake

@@ -108,6 +108,7 @@ endfunction()
 #   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.
 # Outputs:
 #   NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
 function(generate_type_stubs FQ_MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR)
@@ -181,7 +182,7 @@ endfunction()
 function(declare_mlir_python_extension name)
   cmake_parse_arguments(ARG
     ""
-    "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY"
+    "ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY;PACKAGE_PREFIX"
     "SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS;GENERATE_TYPE_STUBS"
     ${ARGN})
 
@@ -207,6 +208,19 @@ function(declare_mlir_python_extension name)
     mlir_python_GENERATE_TYPE_STUBS "${ARG_GENERATE_TYPE_STUBS}"
   )
 
+  if(ARG_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()
+    set(_stubgen_target "${ARG_PACKAGE_PREFIX}._mlir_libs.${name}.type_stub_gen")
+    declare_mlir_python_sources(
+      ${_stubgen_target}
+      ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
+      ADD_TO_PARENT "${ARG_ADD_TO_PARENT}"
+      SOURCES "${ARG_GENERATE_TYPE_STUBS}"
+    )
+  endif()
+
   # Set the interface source and link_libs properties of the target
   # These properties support generator expressions and are automatically exported
   list(TRANSFORM ARG_SOURCES PREPEND "${ARG_ROOT_DIR}/" OUTPUT_VARIABLE _build_sources)
@@ -317,29 +331,15 @@ function(add_mlir_python_modules name)
           message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX")
         endif()
         set(_fully_qualified_module_name "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}")
+        list(TRANSFORM _generate_type_stubs REPLACE "_mlir_libs/" "")
         generate_type_stubs(
           ${_fully_qualified_module_name}
           ${_extension_target}
           ${name}
-          "${CMAKE_CURRENT_BINARY_DIR}/_mlir_libs"
+          "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
           OUTPUTS "${_generate_type_stubs}"
         )
         add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")
-        set(_stubgen_target "${_fully_qualified_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/python/CMakeLists.txt

@@ -40,7 +40,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
     StandaloneCAPI
   PYTHON_BINDINGS_LIBRARY nanobind
   GENERATE_TYPE_STUBS
-    "_standaloneDialectsNanobind.pyi"
+    "_mlir_libs/_standaloneDialectsNanobind.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 
@@ -80,4 +81,4 @@ add_mlir_python_modules(StandalonePythonModules
     StandalonePythonCAPI
   GENERATE_TYPE_STUBS
   PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
-  )
+)

mlir/python/CMakeLists.txt

@@ -506,10 +506,11 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
     # Dialects
     MLIRCAPIFunc
   GENERATE_TYPE_STUBS
-    "_mlir/__init__.pyi"
-    "_mlir/ir.pyi"
-    "_mlir/passmanager.pyi"
-    "_mlir/rewrite.pyi"
+    "_mlir_libs/_mlir/__init__.pyi"
+    "_mlir_libs/_mlir/ir.pyi"
+    "_mlir_libs/_mlir/passmanager.pyi"
+    "_mlir_libs/_mlir/rewrite.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 # This extension exposes an API to register all dialects, extensions, and passes
@@ -532,7 +533,8 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
     MLIRCAPITransforms
     MLIRCAPIRegisterEverything
   GENERATE_TYPE_STUBS
-    "_mlirRegisterEverything.pyi"
+    "_mlir_libs/_mlirRegisterEverything.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -548,7 +550,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
     MLIRCAPIIR
     MLIRCAPILinalg
   GENERATE_TYPE_STUBS
-    "_mlirDialectsLinalg.pyi"
+    "_mlir_libs/_mlirDialectsLinalg.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -564,7 +567,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
     MLIRCAPIIR
     MLIRCAPIGPU
   GENERATE_TYPE_STUBS
-    "_mlirDialectsGPU.pyi"
+    "_mlir_libs/_mlirDialectsGPU.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -580,7 +584,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
     MLIRCAPIIR
     MLIRCAPILLVM
   GENERATE_TYPE_STUBS
-    "_mlirDialectsLLVM.pyi"
+    "_mlir_libs/_mlirDialectsLLVM.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -596,7 +601,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
     MLIRCAPIIR
     MLIRCAPIQuant
   GENERATE_TYPE_STUBS
-    "_mlirDialectsQuant.pyi"
+    "_mlir_libs/_mlirDialectsQuant.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -612,7 +618,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
     MLIRCAPIIR
     MLIRCAPINVGPU
   GENERATE_TYPE_STUBS
-    "_mlirDialectsNVGPU.pyi"
+    "_mlir_libs/_mlirDialectsNVGPU.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -628,7 +635,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
     MLIRCAPIIR
     MLIRCAPIPDL
   GENERATE_TYPE_STUBS
-    "_mlirDialectsPDL.pyi"
+    "_mlir_libs/_mlirDialectsPDL.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -644,7 +652,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
     MLIRCAPIIR
     MLIRCAPISparseTensor
   GENERATE_TYPE_STUBS
-    "_mlirDialectsSparseTensor.pyi"
+    "_mlir_libs/_mlirDialectsSparseTensor.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -660,7 +669,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
     MLIRCAPIIR
     MLIRCAPITransformDialect
   GENERATE_TYPE_STUBS
-    "_mlirDialectsTransform.pyi"
+    "_mlir_libs/_mlirDialectsTransform.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -675,7 +685,8 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIAsync
   GENERATE_TYPE_STUBS
-    "_mlirAsyncPasses.pyi"
+    "_mlir_libs/_mlirAsyncPasses.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 if(MLIR_ENABLE_EXECUTION_ENGINE)
@@ -691,7 +702,8 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIExecutionEngine
     GENERATE_TYPE_STUBS
-      "_mlirExecutionEngine.pyi"
+      "_mlir_libs/_mlirExecutionEngine.pyi"
+    PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
   )
 endif()
 
@@ -707,7 +719,8 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPIGPU
   GENERATE_TYPE_STUBS
-    "_mlirGPUPasses.pyi"
+    "_mlir_libs/_mlirGPUPasses.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -722,7 +735,8 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPILinalg
   GENERATE_TYPE_STUBS
-    "_mlirLinalgPasses.pyi"
+    "_mlir_libs/_mlirLinalgPasses.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -741,7 +755,8 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
     MLIRCAPISMT
     MLIRCAPIExportSMTLIB
   GENERATE_TYPE_STUBS
-    "_mlirDialectsSMT.pyi"
+    "_mlir_libs/_mlirDialectsSMT.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -756,7 +771,8 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
   EMBED_CAPI_LINK_LIBS
     MLIRCAPISparseTensor
   GENERATE_TYPE_STUBS
-    "_mlirSparseTensorPasses.pyi"
+    "_mlir_libs/_mlirSparseTensorPasses.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -771,7 +787,8 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
   EMBED_CAPI_LINK_LIBS
     MLIRCAPITransformDialectTransforms
   GENERATE_TYPE_STUBS
-    "_mlirTransformInterpreter.pyi"
+    "_mlir_libs/_mlirTransformInterpreter.pyi"
+  PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
 )
 
 # TODO: Figure out how to put this in the test tree.
@@ -831,7 +848,8 @@ if(MLIR_INCLUDE_TESTS)
     EMBED_CAPI_LINK_LIBS
       MLIRCAPIPythonTestDialect
     GENERATE_TYPE_STUBS
-      "_mlirPythonTestNanobind.pyi"
+      "_mlir_libs/_mlirPythonTestNanobind.pyi"
+    PACKAGE_PREFIX "${MLIR_PYTHON_PACKAGE_PREFIX}"
   )
 endif()