Skip to content

Commit 4d40913

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

File tree

4 files changed

+61
-62
lines changed

4 files changed

+61
-62
lines changed

‎mlir/cmake/modules/AddMLIRPython.cmake‎

Lines changed: 38 additions & 21 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"
@@ -283,8 +293,8 @@ endfunction()
283293
# extension depend on (see mlir_python_add_common_capi_library).
284294
# GENERATE_TYPE_STUBS: Enable type stub generation for all modules
285295
# 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.
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: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,8 +39,7 @@ declare_mlir_python_extension(StandalonePythonSources.NanobindExtension
3939
EMBED_CAPI_LINK_LIBS
4040
StandaloneCAPI
4141
PYTHON_BINDINGS_LIBRARY nanobind
42-
GENERATE_TYPE_STUBS
43-
"_mlir_libs/_standaloneDialectsNanobind.pyi"
42+
GENERATE_TYPE_STUBS ON
4443
)
4544

4645

‎mlir/python/CMakeLists.txt‎

Lines changed: 21 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,7 @@ declare_mlir_python_extension(MLIRPythonExtension.AsyncDialectPasses
674665
LLVMSupport
675666
EMBED_CAPI_LINK_LIBS
676667
MLIRCAPIAsync
677-
GENERATE_TYPE_STUBS
678-
"_mlir_libs/_mlirAsyncPasses.pyi"
668+
GENERATE_TYPE_STUBS ON
679669
)
680670

681671
if(MLIR_ENABLE_EXECUTION_ENGINE)
@@ -690,8 +680,7 @@ if(MLIR_ENABLE_EXECUTION_ENGINE)
690680
LLVMSupport
691681
EMBED_CAPI_LINK_LIBS
692682
MLIRCAPIExecutionEngine
693-
GENERATE_TYPE_STUBS
694-
"_mlir_libs/_mlirExecutionEngine.pyi"
683+
GENERATE_TYPE_STUBS ON
695684
)
696685
endif()
697686

@@ -706,8 +695,7 @@ declare_mlir_python_extension(MLIRPythonExtension.GPUDialectPasses
706695
LLVMSupport
707696
EMBED_CAPI_LINK_LIBS
708697
MLIRCAPIGPU
709-
GENERATE_TYPE_STUBS
710-
"_mlir_libs/_mlirGPUPasses.pyi"
698+
GENERATE_TYPE_STUBS ON
711699
)
712700

713701
declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
@@ -721,8 +709,7 @@ declare_mlir_python_extension(MLIRPythonExtension.LinalgPasses
721709
LLVMSupport
722710
EMBED_CAPI_LINK_LIBS
723711
MLIRCAPILinalg
724-
GENERATE_TYPE_STUBS
725-
"_mlir_libs/_mlirLinalgPasses.pyi"
712+
GENERATE_TYPE_STUBS ON
726713
)
727714

728715
declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
@@ -740,8 +727,7 @@ declare_mlir_python_extension(MLIRPythonExtension.Dialects.SMT.Pybind
740727
MLIRCAPIIR
741728
MLIRCAPISMT
742729
MLIRCAPIExportSMTLIB
743-
GENERATE_TYPE_STUBS
744-
"_mlir_libs/_mlirDialectsSMT.pyi"
730+
GENERATE_TYPE_STUBS ON
745731
)
746732

747733
declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
@@ -755,8 +741,7 @@ declare_mlir_python_extension(MLIRPythonExtension.SparseTensorDialectPasses
755741
LLVMSupport
756742
EMBED_CAPI_LINK_LIBS
757743
MLIRCAPISparseTensor
758-
GENERATE_TYPE_STUBS
759-
"_mlir_libs/_mlirSparseTensorPasses.pyi"
744+
GENERATE_TYPE_STUBS ON
760745
)
761746

762747
declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
@@ -770,8 +755,7 @@ declare_mlir_python_extension(MLIRPythonExtension.TransformInterpreter
770755
LLVMSupport
771756
EMBED_CAPI_LINK_LIBS
772757
MLIRCAPITransformDialectTransforms
773-
GENERATE_TYPE_STUBS
774-
"_mlir_libs/_mlirTransformInterpreter.pyi"
758+
GENERATE_TYPE_STUBS ON
775759
)
776760

777761
# TODO: Figure out how to put this in the test tree.
@@ -830,8 +814,7 @@ if(MLIR_INCLUDE_TESTS)
830814
LLVMSupport
831815
EMBED_CAPI_LINK_LIBS
832816
MLIRCAPIPythonTestDialect
833-
GENERATE_TYPE_STUBS
834-
"_mlir_libs/_mlirPythonTestNanobind.pyi"
817+
GENERATE_TYPE_STUBS ON
835818
)
836819
endif()
837820

‎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)