Skip to content

Commit d02c7ae

Browse files
maksleventalmakslevental
authored
[MLIR][Python] fix stubgen (llvm#157583)
In llvm#155741 I broke the cardinal rule of CMake: nothing happens when you think it happens 🤷. Meaning: `declare_mlir_python_sources(SOURCES_GLOB "_mlir_libs/${_module_name}/**/*.pyi")` wasn't picking up any sources _because they aren't generated yet_. This of course makes sense in retrospect (the stubs are generated as part of the build process, post extension compile, rather than the configure process). Thus, the API needs to be: ``` GENERATE_TYPE_STUBS: List of generated type stubs expected from stubgen, relative to _mlir_libs. ``` Partially as a result of this omission, the stubs weren't being installed into either the build dir nor the install dir. That is also fixed now: **Source dir (for easy reference):** <img width="300" height="674" alt="image" src="https://github.com/user-attachments/assets/a569f066-c2bd-4361-91f3-1c75181e51da" /> **Build dir (for forthcoming typechecker tests):** <img width="398" height="551" alt="image" src="https://github.com/user-attachments/assets/017859f9-fddb-49ee-85e5-915f5b5f7377" /> **Install dir:** <img width="456" height="884" alt="image" src="https://github.com/user-attachments/assets/8051be7e-898c-4ec8-a11e-e2408b241a56" />
1 parent 7d04e37 commit d02c7ae

File tree

3 files changed

+48
-11
lines changed

3 files changed

+48
-11
lines changed

‎mlir/cmake/modules/AddMLIRPython.cmake‎

Lines changed: 26 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -111,6 +111,11 @@ endfunction()
111111
# Outputs:
112112
# NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
113113
function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR)
114+
cmake_parse_arguments(ARG
115+
""
116+
""
117+
"OUTPUTS"
118+
${ARGN})
114119
if(EXISTS ${nanobind_DIR}/../src/stubgen.py)
115120
set(NB_STUBGEN "${nanobind_DIR}/../src/stubgen.py")
116121
elseif(EXISTS ${nanobind_DIR}/../stubgen.py)
@@ -135,9 +140,9 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP
135140
--output-dir
136141
"${OUTPUT_DIR}")
137142

138-
set(NB_STUBGEN_OUTPUT "${OUTPUT_DIR}/${MODULE_NAME}.pyi")
143+
list(TRANSFORM ARG_OUTPUTS PREPEND "${OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
139144
add_custom_command(
140-
OUTPUT ${NB_STUBGEN_OUTPUT}
145+
OUTPUT ${_generated_type_stubs}
141146
COMMAND ${NB_STUBGEN_CMD}
142147
WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
143148
DEPENDS
@@ -146,7 +151,7 @@ function(generate_type_stubs MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTP
146151
"${DEPENDS_TARGET}"
147152
)
148153
set(_name "MLIRPythonModuleStubs_${_module}")
149-
add_custom_target("${_name}" ALL DEPENDS ${NB_STUBGEN_OUTPUT})
154+
add_custom_target("${_name}" ALL DEPENDS ${_generated_type_stubs})
150155
set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
151156
endfunction()
152157

@@ -166,12 +171,12 @@ endfunction()
166171
# on. These will be collected for all extensions and put into an
167172
# aggregate dylib that is linked against.
168173
# PYTHON_BINDINGS_LIBRARY: Either pybind11 or nanobind.
169-
# GENERATE_TYPE_STUBS: Enable type stub generation.
174+
# GENERATE_TYPE_STUBS: List of generated type stubs expected from stubgen relative to _mlir_libs.
170175
function(declare_mlir_python_extension name)
171176
cmake_parse_arguments(ARG
172-
"GENERATE_TYPE_STUBS"
177+
""
173178
"ROOT_DIR;MODULE_NAME;ADD_TO_PARENT;PYTHON_BINDINGS_LIBRARY"
174-
"SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS"
179+
"SOURCES;PRIVATE_LINK_LIBS;EMBED_CAPI_LINK_LIBS;GENERATE_TYPE_STUBS"
175180
${ARGN})
176181

177182
if(NOT ARG_ROOT_DIR)
@@ -302,15 +307,25 @@ function(add_mlir_python_modules name)
302307
${_module_name}
303308
${_extension_target}
304309
${name}
305-
"${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs/_mlir"
310+
"${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs"
311+
OUTPUTS "${_generate_type_stubs}"
306312
)
313+
add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")
314+
set(_stubgen_target "${MLIR_PYTHON_PACKAGE_PREFIX}.${_module_name}_type_stub_gen")
307315
declare_mlir_python_sources(
308-
"${MLIR_PYTHON_PACKAGE_PREFIX}.${_module_name}_type_stub_gen"
309-
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir"
316+
${_stubgen_target}
317+
ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/mlir/_mlir_libs"
310318
ADD_TO_PARENT "${sources_target}"
311-
SOURCES_GLOB "_mlir_libs/${_module_name}/**/*.pyi"
319+
SOURCES "${_generate_type_stubs}"
312320
)
313-
add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")
321+
set(_pure_sources_target "${modules_target}.sources.${sources_target}_type_stub_gen")
322+
add_mlir_python_sources_target(${_pure_sources_target}
323+
INSTALL_COMPONENT ${modules_target}
324+
INSTALL_DIR "${ARG_INSTALL_PREFIX}/_mlir_libs"
325+
OUTPUT_DIRECTORY "${ARG_ROOT_PREFIX}/_mlir_libs"
326+
SOURCES_TARGETS ${_stubgen_target}
327+
)
328+
add_dependencies(${modules_target} ${_pure_sources_target})
314329
endif()
315330
else()
316331
message(SEND_ERROR "Unrecognized source type '${_source_type}' for python source target ${sources_target}")

‎mlir/python/CMakeLists.txt‎

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -506,6 +506,10 @@ declare_mlir_python_extension(MLIRPythonExtension.Core
506506
# Dialects
507507
MLIRCAPIFunc
508508
GENERATE_TYPE_STUBS
509+
"_mlir/__init__.pyi"
510+
"_mlir/ir.pyi"
511+
"_mlir/passmanager.pyi"
512+
"_mlir/rewrite.pyi"
509513
)
510514

511515
# This extension exposes an API to register all dialects, extensions, and passes
@@ -528,6 +532,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
528532
MLIRCAPITransforms
529533
MLIRCAPIRegisterEverything
530534
GENERATE_TYPE_STUBS
535+
"_mlirRegisterEverything.pyi"
531536
)
532537

533538
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -543,6 +548,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
543548
MLIRCAPIIR
544549
MLIRCAPILinalg
545550
GENERATE_TYPE_STUBS
551+
"_mlirDialectsLinalg.pyi"
546552
)
547553

548554
declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -558,6 +564,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
558564
MLIRCAPIIR
559565
MLIRCAPIGPU
560566
GENERATE_TYPE_STUBS
567+
"_mlirDialectsGPU.pyi"
561568
)
562569

563570
declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -573,6 +580,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
573580
MLIRCAPIIR
574581
MLIRCAPILLVM
575582
GENERATE_TYPE_STUBS
583+
"_mlirDialectsLLVM.pyi"
576584
)
577585

578586
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -588,6 +596,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
588596
MLIRCAPIIR
589597
MLIRCAPIQuant
590598
GENERATE_TYPE_STUBS
599+
"_mlirDialectsQuant.pyi"
591600
)
592601

593602
declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -603,6 +612,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
603612
MLIRCAPIIR
604613
MLIRCAPINVGPU
605614
GENERATE_TYPE_STUBS
615+
"_mlirDialectsNVGPU.pyi"
606616
)
607617

608618
declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -618,6 +628,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
618628
MLIRCAPIIR
619629
MLIRCAPIPDL
620630
GENERATE_TYPE_STUBS
631+
"_mlirDialectsPDL.pyi"
621632
)
622633

623634
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -633,6 +644,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
633644
MLIRCAPIIR
634645
MLIRCAPISparseTensor
635646
GENERATE_TYPE_STUBS
647+
"_mlirDialectsSparseTensor.pyi"
636648
)
637649

638650
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -648,6 +660,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
648660
MLIRCAPIIR
649661
MLIRCAPITransformDialect
650662
GENERATE_TYPE_STUBS
663+
"_mlirDialectsTransform.pyi"
651664
)
652665

653666
declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -662,6 +675,7 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
662675
EMBED_CAPI_LINK_LIBS
663676
MLIRCAPIAsync
664677
GENERATE_TYPE_STUBS
678+
"_mlirAsyncPasses.pyi"
665679
)
666680

667681
if(MLIR_ENABLE_EXECUTION_ENGINE)
@@ -677,6 +691,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
677691
EMBED_CAPI_LINK_LIBS
678692
MLIRCAPIExecutionEngine
679693
GENERATE_TYPE_STUBS
694+
"_mlirExecutionEngine.pyi"
680695
)
681696
endif()
682697

@@ -692,6 +707,7 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
692707
EMBED_CAPI_LINK_LIBS
693708
MLIRCAPIGPU
694709
GENERATE_TYPE_STUBS
710+
"_mlirGPUPasses.pyi"
695711
)
696712

697713
declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -706,6 +722,7 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
706722
EMBED_CAPI_LINK_LIBS
707723
MLIRCAPILinalg
708724
GENERATE_TYPE_STUBS
725+
"_mlirLinalgPasses.pyi"
709726
)
710727

711728
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -724,6 +741,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
724741
MLIRCAPISMT
725742
MLIRCAPIExportSMTLIB
726743
GENERATE_TYPE_STUBS
744+
"_mlirDialectsSMT.pyi"
727745
)
728746

729747
declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -738,6 +756,7 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
738756
EMBED_CAPI_LINK_LIBS
739757
MLIRCAPISparseTensor
740758
GENERATE_TYPE_STUBS
759+
"_mlirSparseTensorPasses.pyi"
741760
)
742761

743762
declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -752,6 +771,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
752771
EMBED_CAPI_LINK_LIBS
753772
MLIRCAPITransformDialectTransforms
754773
GENERATE_TYPE_STUBS
774+
"_mlirTransformInterpreter.pyi"
755775
)
756776

757777
# TODO: Figure out how to put this in the test tree.
@@ -811,6 +831,7 @@ if(MLIR_INCLUDE_TESTS)
811831
EMBED_CAPI_LINK_LIBS
812832
MLIRCAPIPythonTestDialect
813833
GENERATE_TYPE_STUBS
834+
"_mlirPythonTestNanobind.pyi"
814835
)
815836
endif()
816837

‎mlir/python/mlir/_mlir_libs/.gitignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,2 @@
11
_mlir/**/*.pyi
2+
*.pyi

0 commit comments

Comments
 (0)