Skip to content

Commit 3e13b03

Browse files
tlopextlopex
authored
[Build] Migrate Python packaging to pyproject.toml with scikit-build-core (#18239)
This pr migrates the TVM Python packaging system from the setup.py flow to the modern, PEP 517/518 compliant pyproject.toml standard, which allows us to produce a single, Python-version-agnostic wheel. This change streamlines the process for both developers and users. For local development, you can now set up a fully-functional editable environment with a single command: `pip install -e .`. To create the distributable package for release, simply run `pip wheel -w dist .` , which will produce a universal wheel in the `dist/` folder. This ensures that end-users can reliably install TVM with a standard pip install tvm, regardless of their specific Python 3 version.
1 parent dd1e3f8 commit 3e13b03

File tree

4 files changed

+320
-1
lines changed

4 files changed

+320
-1
lines changed

‎CMakeLists.txt‎

Lines changed: 90 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,9 @@ tvm_option(USE_MSC "Enable Multi-System Compiler" OFF)
121121
tvm_option(USE_MRVL "Build with MRVL TVM support" OFF)
122122
tvm_option(USE_NVSHMEM "Build with NVSHMEM support" OFF)
123123

124+
# Python package options
125+
tvm_option(TVM_BUILD_PYTHON_MODULE "Build Python module with scikit-build-core" ON)
126+
124127
# include directories
125128
include_directories(${CMAKE_INCLUDE_PATH})
126129
include_directories("include")
@@ -566,7 +569,6 @@ endif()
566569

567570
add_subdirectory(ffi)
568571

569-
570572
if(TVM_DEBUG_WITH_ABI_CHANGE)
571573
message(STATUS "Building with debug code that may cause ABI changes...")
572574
target_compile_definitions(tvm_objs PRIVATE "TVM_DEBUG_WITH_ABI_CHANGE")
@@ -818,3 +820,90 @@ if(USE_ROCM AND USE_RCCL)
818820
target_link_libraries(tvm PRIVATE rccl)
819821
target_link_libraries(tvm_runtime PRIVATE rccl)
820822
endif()
823+
824+
# Python package installation configuration
825+
# This section ensures that all necessary files are installed for the Python wheel
826+
if(TVM_BUILD_PYTHON_MODULE)
827+
message(STATUS "Configuring Python package installation")
828+
829+
# Install Python source files
830+
install(
831+
DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/python/tvm"
832+
DESTINATION "."
833+
FILES_MATCHING
834+
PATTERN "*.py"
835+
PATTERN "*.pyi"
836+
PATTERN "__pycache__" EXCLUDE
837+
PATTERN "*.pyc" EXCLUDE
838+
)
839+
840+
# Install compiled shared libraries
841+
install(TARGETS tvm DESTINATION "tvm")
842+
install(TARGETS tvm_runtime DESTINATION "tvm")
843+
844+
# Install third-party compiled dependencies
845+
if(TARGET fpA_intB_gemm)
846+
install(TARGETS fpA_intB_gemm DESTINATION "tvm")
847+
endif()
848+
if(TARGET flash_attn)
849+
install(TARGETS flash_attn DESTINATION "tvm")
850+
endif()
851+
852+
# Install minimal header files needed by Python extensions
853+
install(
854+
DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/include/tvm/runtime"
855+
DESTINATION "tvm/include/tvm/runtime"
856+
FILES_MATCHING
857+
PATTERN "*.h"
858+
)
859+
860+
# Install minimal CMake configuration
861+
install(
862+
DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/cmake/utils"
863+
DESTINATION "tvm/cmake/utils"
864+
FILES_MATCHING
865+
PATTERN "*.cmake"
866+
)
867+
868+
# Install CUTLASS headers only if available
869+
if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/cutlass/include")
870+
install(
871+
DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/3rdparty/cutlass/include"
872+
DESTINATION "tvm/3rdparty/cutlass"
873+
FILES_MATCHING
874+
PATTERN "*.h"
875+
PATTERN "*.hpp"
876+
)
877+
endif()
878+
879+
# Install minimal source files
880+
install(
881+
DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/src/runtime"
882+
DESTINATION "tvm/src/runtime"
883+
FILES_MATCHING
884+
PATTERN "*.cc"
885+
PATTERN "*.h"
886+
)
887+
888+
# Install essential configuration files
889+
install(
890+
DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/configs"
891+
DESTINATION "tvm/configs"
892+
)
893+
894+
# Install licenses (required for distribution)
895+
install(
896+
DIRECTORY "${CMAKE_CURRENT_SOURCE_DIR}/licenses"
897+
DESTINATION "tvm/licenses"
898+
)
899+
900+
# Install essential metadata files
901+
install(FILES
902+
"${CMAKE_CURRENT_SOURCE_DIR}/README.md"
903+
"${CMAKE_CURRENT_SOURCE_DIR}/LICENSE"
904+
"${CMAKE_CURRENT_SOURCE_DIR}/NOTICE"
905+
DESTINATION "tvm"
906+
)
907+
908+
message(STATUS "Python package installation configured")
909+
endif()

‎cmake/modules/LibInfo.cmake‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ function(add_lib_info src_file)
107107
TVM_INFO_USE_ROCM="${USE_ROCM}"
108108
TVM_INFO_USE_RCCL="${USE_RCCL}"
109109
TVM_INFO_USE_RPC="${USE_RPC}"
110+
TVM_INFO_TVM_BUILD_PYTHON_MODULE="${TVM_BUILD_PYTHON_MODULE}"
110111
TVM_INFO_USE_RTTI="${USE_RTTI}"
111112
TVM_INFO_USE_RUST_EXT="${USE_RUST_EXT}"
112113
TVM_INFO_USE_SORT="${USE_SORT}"

‎pyproject.toml‎

Lines changed: 228 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,191 @@
1414
# KIND, either express or implied. See the License for the
1515
# specific language governing permissions and limitations
1616
# under the License.
17+
18+
[build-system]
19+
requires = ["scikit-build-core>=0.10.0"]
20+
build-backend = "scikit_build_core.build"
21+
22+
[project]
23+
name = "tvm"
24+
# Note: Call version.py to update the version before building the wheel
25+
version = "0.22.0.dev0"
26+
description = "Apache TVM: An End-to-End Deep Learning Compiler Stack"
27+
readme = "README.md"
28+
license = { text = "Apache-2.0" }
29+
requires-python = ">=3.9"
30+
authors = [
31+
{ name = "Apache TVM Community", email = "[email protected]" }
32+
]
33+
keywords = ["machine learning", "compiler", "deep learning", "inference"]
34+
classifiers = [
35+
"Development Status :: 4 - Beta",
36+
"Intended Audience :: Developers",
37+
"Intended Audience :: Education",
38+
"Intended Audience :: Science/Research",
39+
"License :: OSI Approved :: Apache Software License",
40+
"Programming Language :: Python :: 3",
41+
"Programming Language :: Python :: 3.9",
42+
"Programming Language :: Python :: 3.10",
43+
"Programming Language :: Python :: 3.11",
44+
"Programming Language :: Python :: 3.12",
45+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
46+
"Topic :: Software Development :: Libraries :: Python Modules",
47+
]
48+
# Core dependencies - these are the minimum required for basic TVM functionality
49+
dependencies = [
50+
"cloudpickle",
51+
"ml_dtypes",
52+
"numpy",
53+
"packaging",
54+
"psutil",
55+
"scipy",
56+
"tornado",
57+
"typing_extensions",
58+
]
59+
60+
# Optional dependencies for different features
61+
[project.optional-dependencies]
62+
# Model importers
63+
importer-coreml = ["coremltools"]
64+
importer-keras = ["tensorflow", "tensorflow-estimator"]
65+
importer-onnx = ["future", "onnx", "onnxoptimizer", "onnxruntime", "torch", "torchvision"]
66+
importer-pytorch = ["torch", "torchvision"]
67+
importer-tensorflow = ["tensorflow", "tensorflow-estimator"]
68+
importer-tflite = ["tflite"]
69+
importer-paddle = ["paddlepaddle"]
70+
71+
# AutoTVM and autoscheduler
72+
autotvm = ["xgboost"]
73+
autoscheduler = ["xgboost"]
74+
75+
# Development and testing
76+
dev = [
77+
"black",
78+
"isort",
79+
"mypy",
80+
"pylint",
81+
"pytest",
82+
"pytest-xdist",
83+
"pytest-cov",
84+
"pytest-mock",
85+
"pytest-benchmark",
86+
"pytest-timeout",
87+
"pytest-rerunfailures",
88+
"pytest-repeat",
89+
"pytest-xdist",
90+
"pytest-cov",
91+
"pytest-mock",
92+
"pytest-benchmark",
93+
"pytest-timeout",
94+
"pytest-rerunfailures",
95+
"pytest-repeat",
96+
]
97+
98+
# All optional dependencies (excluding dev)
99+
all = [
100+
"coremltools",
101+
"tensorflow",
102+
"tensorflow-estimator",
103+
"future",
104+
"onnx",
105+
"onnxoptimizer",
106+
"onnxruntime",
107+
"torch",
108+
"torchvision",
109+
"tflite",
110+
"paddlepaddle",
111+
"xgboost",
112+
]
113+
114+
[project.urls]
115+
Homepage = "https://tvm.apache.org/"
116+
Documentation = "https://tvm.apache.org/docs/"
117+
Repository = "https://github.com/apache/tvm"
118+
"Bug Tracker" = "https://github.com/apache/tvm/issues"
119+
120+
[tool.scikit-build]
121+
# Point to the root CMakeLists.txt
122+
cmake.source-dir = "."
123+
cmake.build-type = "Release"
124+
125+
# Configure the wheel to be Python version-agnostic
126+
wheel.py-api = "py3"
127+
128+
# Build configuration
129+
build-dir = "build"
130+
131+
# CMake configuration - ensure proper installation paths
132+
cmake.args = [
133+
"-DTVM_BUILD_PYTHON_MODULE=ON",
134+
"-DTVM_FFI_BUILD_PYTHON_MODULE=OFF",
135+
"-DTVM_USE_CUTLASS=OFF",
136+
"-DTVM_USE_FLASH_ATTN=OFF",
137+
"-DTVM_USE_LLVM=OFF",
138+
"-DTVM_USE_CUDA=OFF",
139+
"-DTVM_USE_OPENCL=OFF",
140+
"-DTVM_USE_VULKAN=OFF",
141+
"-DTVM_USE_METAL=OFF",
142+
"-DTVM_USE_OPENGL=OFF",
143+
"-DTVM_USE_RPC=OFF",
144+
"-DTVM_USE_GRAPH_EXECUTOR=OFF",
145+
"-DTVM_USE_PROFILER=OFF",
146+
"-DTVM_USE_UTILS=OFF",
147+
]
148+
149+
# Wheel configuration
150+
wheel.packages = ["python/tvm"]
151+
152+
# Source distribution configuration
153+
sdist.include = [
154+
# Build files
155+
"/CMakeLists.txt",
156+
"/pyproject.toml",
157+
"/cmake/**/*",
158+
"/3rdparty/**/*",
159+
160+
# Source code
161+
"/src/**/*.cc",
162+
"/src/**/*.h",
163+
"/include/**/*.h",
164+
165+
# Python source
166+
"/python/tvm/**/*.py",
167+
"/python/tvm/**/*.pyi",
168+
169+
# Documentation and metadata
170+
"/docs/**/*",
171+
"/LICENSE",
172+
"/README.md",
173+
"/NOTICE",
174+
175+
# Tests
176+
"/tests/**/*",
177+
]
178+
179+
sdist.exclude = [
180+
"**/.git",
181+
"**/.github",
182+
"**/__pycache__",
183+
"**/*.pyc",
184+
"build",
185+
"dist",
186+
"**/3rdparty/*/docs",
187+
"**/3rdparty/*/media",
188+
"**/3rdparty/*/examples",
189+
"**/3rdparty/*/test",
190+
]
191+
192+
# Logging
193+
logging.level = "INFO"
194+
195+
[tool.pytest.ini_options]
196+
testpaths = ["tests"]
197+
addopts = "-v --tb=short"
198+
python_files = ["test_*.py", "*_test.py"]
199+
python_classes = ["Test*"]
200+
python_functions = ["test_*"]
201+
17202
[tool.isort]
18203
profile = "black"
19204
src_paths = ["python", "tests/python"]
@@ -51,5 +236,48 @@ exclude = '''
51236
'''
52237

53238
[tool.ruff]
239+
# Enable pycodestyle (`E`), Pyflakes (`F`), and isort (`I`) codes
240+
select = ["E", "F", "I"]
241+
ignore = []
242+
243+
# Allow fix for all enabled rules (when `--fix`) is provided.
244+
fixable = ["A", "B", "C", "D", "E", "F", "I", "N", "UP", "W", "ARG", "B", "C4", "DTZ", "T10", "EM", "EXE", "FA", "ICN", "Q", "T20", "TID", "TCH", "RUF"]
245+
unfixable = []
246+
247+
# Exclude a variety of commonly ignored directories.
248+
exclude = [
249+
".bzr",
250+
".darcs",
251+
".git",
252+
".git",
253+
".hg",
254+
".mypy_cache",
255+
".nox",
256+
".pants.d",
257+
".pytype",
258+
".ruff_cache",
259+
".svn",
260+
".tox",
261+
".venv",
262+
"__pypackages__",
263+
"_build",
264+
"buck-out",
265+
"build",
266+
"dist",
267+
"node_modules",
268+
"venv",
269+
"3rdparty",
270+
]
271+
54272
line-length = 100
55273
indent-width = 4
274+
275+
# Allow unused variables when underscore-prefixed.
276+
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
277+
278+
279+
[tool.ruff.mccabe]
280+
max-complexity = 10
281+
282+
[tool.ruff.isort]
283+
known-first-party = ["tvm"]

‎src/support/libinfo.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -339,6 +339,7 @@ TVM_DLL ffi::Map<ffi::String, ffi::String> GetLibInfo() {
339339
{"USE_ROCM", TVM_INFO_USE_ROCM},
340340
{"USE_RCCL", TVM_INFO_USE_RCCL},
341341
{"USE_RPC", TVM_INFO_USE_RPC},
342+
{"TVM_BUILD_PYTHON_MODULE", TVM_INFO_TVM_BUILD_PYTHON_MODULE},
342343
{"USE_RTTI", TVM_INFO_USE_RTTI},
343344
{"USE_RUST_EXT", TVM_INFO_USE_RUST_EXT},
344345
{"USE_SORT", TVM_INFO_USE_SORT},

0 commit comments

Comments
 (0)