Skip to content

Commit 2483348

Browse files
maksleventalmakslevental
committed
remove mlir_libs prefix
1 parent 2ba888e commit 2483348

File tree

4 files changed

+59
-62
lines changed

4 files changed

+59
-62
lines changed

‎mlir/cmake/modules/AddMLIRPython.cmake‎

Lines changed: 39 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -105,16 +105,16 @@ endfunction()
105105
# FQ_MODULE_NAME: The fully-qualified name of the extension module (used for importing in python).
106106
# DEPENDS_TARGET: The dso target corresponding to the extension module
107107
# (e.g., something like StandalonePythonModules.extension._standaloneDialectsNanobind.dso)
108-
# MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module
108+
# CORE_MLIR_DEPENDS_TARGET: The dso target corresponding to the main/core extension module
109109
# (e.g., something like StandalonePythonModules.extension._mlir.dso)
110110
# OUTPUT_DIR: The root output directory to emit the type stubs into.
111111
# OUTPUTS: List of expected outputs.
112112
# Outputs:
113113
# NB_STUBGEN_CUSTOM_TARGET: The target corresponding to generation which other targets can depend on.
114-
function(generate_type_stubs FQ_MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET OUTPUT_DIR)
114+
function(generate_type_stubs )
115115
cmake_parse_arguments(ARG
116116
""
117-
""
117+
"FQ_MODULE_NAME;DEPENDS_TARGET;CORE_MLIR_DEPENDS_TARGET;OUTPUT_DIR"
118118
"OUTPUTS"
119119
${ARGN})
120120
# for people doing find_package(nanobind)
@@ -134,30 +134,30 @@ function(generate_type_stubs FQ_MODULE_NAME DEPENDS_TARGET MLIR_DEPENDS_TARGET O
134134

135135
file(REAL_PATH "${MLIR_BINARY_DIR}/${MLIR_BINDINGS_PYTHON_INSTALL_PREFIX}/.." _import_path)
136136

137-
set(NB_STUBGEN_CMD
137+
set(_nb_stubgen_cmd
138138
"${Python_EXECUTABLE}"
139139
"${NB_STUBGEN}"
140140
--module
141-
"${FQ_MODULE_NAME}"
141+
"${ARG_FQ_MODULE_NAME}"
142142
-i
143143
"${_import_path}"
144144
--recursive
145145
--include-private
146146
--output-dir
147-
"${OUTPUT_DIR}"
147+
"${ARG_OUTPUT_DIR}"
148148
--quiet)
149149

150-
list(TRANSFORM ARG_OUTPUTS PREPEND "${OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
150+
list(TRANSFORM ARG_OUTPUTS PREPEND "${ARG_OUTPUT_DIR}/" OUTPUT_VARIABLE _generated_type_stubs)
151151
add_custom_command(
152152
OUTPUT ${_generated_type_stubs}
153-
COMMAND ${NB_STUBGEN_CMD}
153+
COMMAND ${_nb_stubgen_cmd}
154154
WORKING_DIRECTORY "${CMAKE_CURRENT_FUNCTION_LIST_DIR}"
155155
DEPENDS
156-
"${MLIR_DEPENDS_TARGET}.extension._mlir.dso"
157-
"${MLIR_DEPENDS_TARGET}.sources.MLIRPythonSources.Core.Python"
158-
"${DEPENDS_TARGET}"
156+
"${ARG_CORE_MLIR_DEPENDS_TARGET}.extension._mlir.dso"
157+
"${ARG_CORE_MLIR_DEPENDS_TARGET}.sources.MLIRPythonSources.Core.Python"
158+
"${ARG_DEPENDS_TARGET}"
159159
)
160-
set(_name "${FQ_MODULE_NAME}.type_stubs")
160+
set(_name "${ARG_FQ_MODULE_NAME}.type_stubs")
161161
add_custom_target("${_name}" DEPENDS ${_generated_type_stubs})
162162
set(NB_STUBGEN_CUSTOM_TARGET "${_name}" PARENT_SCOPE)
163163
endfunction()
@@ -178,7 +178,10 @@ endfunction()
178178
# on. These will be collected for all extensions and put into an
179179
# aggregate dylib that is linked against.
180180
# PYTHON_BINDINGS_LIBRARY: Either pybind11 or nanobind.
181-
# GENERATE_TYPE_STUBS: List of generated type stubs expected from stubgen relative to _mlir_libs.
181+
# GENERATE_TYPE_STUBS: Either
182+
# 1. OFF (default)
183+
# 2. ON if ${MODULE_NAME}.pyi is the only stub
184+
# 3. A list of generated type stubs expected from stubgen relative to _mlir_libs.
182185
function(declare_mlir_python_extension name)
183186
cmake_parse_arguments(ARG
184187
""
@@ -197,7 +200,14 @@ function(declare_mlir_python_extension name)
197200

198201
add_library(${name} INTERFACE)
199202

203+
if(NOT ARG_GENERATE_TYPE_STUBS)
204+
set(ARG_GENERATE_TYPE_STUBS OFF)
205+
endif()
206+
if("${ARG_GENERATE_TYPE_STUBS}" STREQUAL "ON")
207+
set(ARG_GENERATE_TYPE_STUBS "${ARG_MODULE_NAME}.pyi")
208+
endif()
200209
if(ARG_GENERATE_TYPE_STUBS)
210+
list(TRANSFORM ARG_GENERATE_TYPE_STUBS PREPEND "_mlir_libs/")
201211
declare_mlir_python_sources(
202212
"${name}.type_stub_gen"
203213
ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs"
@@ -282,9 +292,9 @@ endfunction()
282292
# COMMON_CAPI_LINK_LIBS: List of dylibs (typically one) to make every
283293
# extension depend on (see mlir_python_add_common_capi_library).
284294
# GENERATE_TYPE_STUBS: Enable type stub generation for all modules
285-
# which have enabled type stub generation.
286-
# PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX. This is used
287-
# to determine type stub generation python module names.
295+
# which have declared type stub generation with `declare_mlir_python_extension`.
296+
# PACKAGE_PREFIX: Same as MLIR_PYTHON_PACKAGE_PREFIX not including trailing `.`.
297+
# This is used to determine type stub generation python module names.
288298
function(add_mlir_python_modules name)
289299
cmake_parse_arguments(ARG
290300
"GENERATE_TYPE_STUBS"
@@ -325,15 +335,22 @@ function(add_mlir_python_modules name)
325335
get_target_property(_generate_type_stubs ${sources_target} mlir_python_GENERATE_TYPE_STUBS)
326336
if(ARG_GENERATE_TYPE_STUBS AND _generate_type_stubs)
327337
if ((NOT ARG_PACKAGE_PREFIX) OR ("${ARG_PACKAGE_PREFIX}" STREQUAL ""))
328-
message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX")
338+
message(FATAL_ERROR "GENERATE_TYPE_STUBS requires PACKAGE_PREFIX for ${name}")
329339
endif()
330-
set(_fully_qualified_module_name "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}")
340+
# TL;DR: everything here is load bearing and annoyingly coupled. Changing anything here
341+
# (or in declare_mlir_python_extension) will break either in-tree or out-of-tree generation.
342+
#
343+
# We remove _mlir_libs here because OUTPUT_DIR already includes it.
344+
# Specifically OUTPUT_DIR already includes it because that's the actual directory
345+
# where we want stubgen to dump the emitted sources. This is load bearing because up above
346+
# (in declare_mlir_python_extension) we prefixed all the paths with _mlir_libs and
347+
# we specified declare_mlir_python_sources with ROOT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs".
331348
list(TRANSFORM _generate_type_stubs REPLACE "_mlir_libs/" "")
332349
generate_type_stubs(
333-
${_fully_qualified_module_name}
334-
${_extension_target}
335-
${name}
336-
"${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
350+
FQ_MODULE_NAME "${ARG_PACKAGE_PREFIX}._mlir_libs.${_module_name}"
351+
DEPENDS_TARGET ${_extension_target}
352+
CORE_MLIR_DEPENDS_TARGET ${name}
353+
OUTPUT_DIR "${CMAKE_CURRENT_BINARY_DIR}/type_stubs/_mlir_libs"
337354
OUTPUTS "${_generate_type_stubs}"
338355
)
339356
add_dependencies("${modules_target}" "${NB_STUBGEN_CUSTOM_TARGET}")

‎mlir/examples/standalone/python/CMakeLists.txt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -40,7 +40,8 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
4040
StandaloneCAPI
4141
PYTHON_BINDINGS_LIBRARY nanobind
4242
GENERATE_TYPE_STUBS
43-
"_mlir_libs/_standaloneDialectsNanobind.pyi"
43+
"_standaloneDialectsNanobind/__init__.pyi"
44+
"_standaloneDialectsNanobind/standalone.pyi"
4445
)
4546

4647

‎mlir/python/CMakeLists.txt‎

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

515515
# This extension exposes an API to register all dialects, extensions, and passes
@@ -531,8 +531,7 @@ declare_mlir_python_extension(MLIRPythonExtension.RegisterEverything
531531
MLIRCAPIConversion
532532
MLIRCAPITransforms
533533
MLIRCAPIRegisterEverything
534-
GENERATE_TYPE_STUBS
535-
"_mlir_libs/_mlirRegisterEverything.pyi"
534+
GENERATE_TYPE_STUBS ON
536535
)
537536

538537
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
@@ -547,8 +546,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Linalg.Pybind
547546
EMBED_CAPI_LINK_LIBS
548547
MLIRCAPIIR
549548
MLIRCAPILinalg
550-
GENERATE_TYPE_STUBS
551-
"_mlir_libs/_mlirDialectsLinalg.pyi"
549+
GENERATE_TYPE_STUBS ON
552550
)
553551

554552
declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
@@ -563,8 +561,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.GPU.Pybind
563561
EMBED_CAPI_LINK_LIBS
564562
MLIRCAPIIR
565563
MLIRCAPIGPU
566-
GENERATE_TYPE_STUBS
567-
"_mlir_libs/_mlirDialectsGPU.pyi"
564+
GENERATE_TYPE_STUBS ON
568565
)
569566

570567
declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
@@ -579,8 +576,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.LLVM.Pybind
579576
EMBED_CAPI_LINK_LIBS
580577
MLIRCAPIIR
581578
MLIRCAPILLVM
582-
GENERATE_TYPE_STUBS
583-
"_mlir_libs/_mlirDialectsLLVM.pyi"
579+
GENERATE_TYPE_STUBS ON
584580
)
585581

586582
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
@@ -595,8 +591,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Quant.Pybind
595591
EMBED_CAPI_LINK_LIBS
596592
MLIRCAPIIR
597593
MLIRCAPIQuant
598-
GENERATE_TYPE_STUBS
599-
"_mlir_libs/_mlirDialectsQuant.pyi"
594+
GENERATE_TYPE_STUBS ON
600595
)
601596

602597
declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
@@ -611,8 +606,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.NVGPU.Pybind
611606
EMBED_CAPI_LINK_LIBS
612607
MLIRCAPIIR
613608
MLIRCAPINVGPU
614-
GENERATE_TYPE_STUBS
615-
"_mlir_libs/_mlirDialectsNVGPU.pyi"
609+
GENERATE_TYPE_STUBS ON
616610
)
617611

618612
declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
@@ -627,8 +621,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.PDL.Pybind
627621
EMBED_CAPI_LINK_LIBS
628622
MLIRCAPIIR
629623
MLIRCAPIPDL
630-
GENERATE_TYPE_STUBS
631-
"_mlir_libs/_mlirDialectsPDL.pyi"
624+
GENERATE_TYPE_STUBS ON
632625
)
633626

634627
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
@@ -643,8 +636,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SparseTensor.Pybind
643636
EMBED_CAPI_LINK_LIBS
644637
MLIRCAPIIR
645638
MLIRCAPISparseTensor
646-
GENERATE_TYPE_STUBS
647-
"_mlir_libs/_mlirDialectsSparseTensor.pyi"
639+
GENERATE_TYPE_STUBS ON
648640
)
649641

650642
declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
@@ -659,8 +651,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.Transform.Pybind
659651
EMBED_CAPI_LINK_LIBS
660652
MLIRCAPIIR
661653
MLIRCAPITransformDialect
662-
GENERATE_TYPE_STUBS
663-
"_mlir_libs/_mlirDialectsTransform.pyi"
654+
GENERATE_TYPE_STUBS ON
664655
)
665656

666657
declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
@@ -674,8 +665,6 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
674665
LLVMSupport
675666
EMBED_CAPI_LINK_LIBS
676667
MLIRCAPIAsync
677-
GENERATE_TYPE_STUBS
678-
"_mlir_libs/_mlirAsyncPasses.pyi"
679668
)
680669

681670
if(MLIR_ENABLE_EXECUTION_ENGINE)
@@ -690,8 +679,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
690679
LLVMSupport
691680
EMBED_CAPI_LINK_LIBS
692681
MLIRCAPIExecutionEngine
693-
GENERATE_TYPE_STUBS
694-
"_mlir_libs/_mlirExecutionEngine.pyi"
682+
GENERATE_TYPE_STUBS ON
695683
)
696684
endif()
697685

@@ -706,8 +694,6 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
706694
LLVMSupport
707695
EMBED_CAPI_LINK_LIBS
708696
MLIRCAPIGPU
709-
GENERATE_TYPE_STUBS
710-
"_mlir_libs/_mlirGPUPasses.pyi"
711697
)
712698

713699
declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -721,8 +707,6 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
721707
LLVMSupport
722708
EMBED_CAPI_LINK_LIBS
723709
MLIRCAPILinalg
724-
GENERATE_TYPE_STUBS
725-
"_mlir_libs/_mlirLinalgPasses.pyi"
726710
)
727711

728712
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -740,8 +724,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
740724
MLIRCAPIIR
741725
MLIRCAPISMT
742726
MLIRCAPIExportSMTLIB
743-
GENERATE_TYPE_STUBS
744-
"_mlir_libs/_mlirDialectsSMT.pyi"
727+
GENERATE_TYPE_STUBS ON
745728
)
746729

747730
declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -755,8 +738,6 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
755738
LLVMSupport
756739
EMBED_CAPI_LINK_LIBS
757740
MLIRCAPISparseTensor
758-
GENERATE_TYPE_STUBS
759-
"_mlir_libs/_mlirSparseTensorPasses.pyi"
760741
)
761742

762743
declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -770,8 +751,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
770751
LLVMSupport
771752
EMBED_CAPI_LINK_LIBS
772753
MLIRCAPITransformDialectTransforms
773-
GENERATE_TYPE_STUBS
774-
"_mlir_libs/_mlirTransformInterpreter.pyi"
754+
GENERATE_TYPE_STUBS ON
775755
)
776756

777757
# TODO: Figure out how to put this in the test tree.
@@ -830,8 +810,7 @@ if(MLIR_INCLUDE_TESTS)
830810
LLVMSupport
831811
EMBED_CAPI_LINK_LIBS
832812
MLIRCAPIPythonTestDialect
833-
GENERATE_TYPE_STUBS
834-
"_mlir_libs/_mlirPythonTestNanobind.pyi"
813+
GENERATE_TYPE_STUBS ON
835814
)
836815
endif()
837816

‎mlir/test/python/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ add_subdirectory(lib)
1313

1414
set(MLIR_PYTHON_TEST_DEPENDS MLIRPythonModules)
1515
if(NOT MLIR_STANDALONE_BUILD)
16-
list(APPEND MLIR_PYTHON_TEST_DEPENDS FileCheck count)
16+
list(APPEND MLIR_PYTHON_TEST_DEPENDS FileCheck count not split-file yaml2obj)
1717
endif()
1818
add_lit_testsuite(check-mlir-python "Running the MLIR Python regression tests"
1919
${CMAKE_CURRENT_BINARY_DIR}

0 commit comments

Comments
 (0)