From 2fdec42d971bbae43d52fa366b9a487c024dfe92 Mon Sep 17 00:00:00 2001 From: tqchen Date: Sat, 30 Aug 2025 16:34:14 -0400 Subject: [PATCH] [FFI] Wheel packaging example This PR add an example about wheel packaging. Also fixes various source packaging minor nits. --- ffi/CMakeLists.txt | 4 +- ffi/cmake/tvm_ffi-config.cmake | 2 + ffi/examples/get_started/README.md | 7 +- ffi/examples/packaging/CMakeLists.txt | 73 +++++++++++++++++++ ffi/examples/packaging/README.md | 61 ++++++++++++++++ ffi/examples/packaging/pyproject.toml | 58 +++++++++++++++ .../python/tvm_ffi_extension/__init__.py | 48 ++++++++++++ .../python/tvm_ffi_extension/_ffi_api.py | 24 ++++++ .../python/tvm_ffi_extension/base.py | 37 ++++++++++ ffi/pyproject.toml | 2 +- ffi/python/tvm_ffi/cython/function.pxi | 7 +- ffi/python/tvm_ffi/libinfo.py | 2 +- 12 files changed, 313 insertions(+), 12 deletions(-) create mode 100644 ffi/examples/packaging/CMakeLists.txt create mode 100644 ffi/examples/packaging/README.md create mode 100644 ffi/examples/packaging/pyproject.toml create mode 100644 ffi/examples/packaging/python/tvm_ffi_extension/__init__.py create mode 100644 ffi/examples/packaging/python/tvm_ffi_extension/_ffi_api.py create mode 100644 ffi/examples/packaging/python/tvm_ffi_extension/base.py diff --git a/ffi/CMakeLists.txt b/ffi/CMakeLists.txt index f40313636ac8..90f1f89cbb92 100644 --- a/ffi/CMakeLists.txt +++ b/ffi/CMakeLists.txt @@ -239,9 +239,9 @@ if (TVM_FFI_BUILD_PYTHON_MODULE) PATTERN "*.tmp" EXCLUDE ) install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/ DESTINATION src/ffi/) - install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/cmake/Utils/ DESTINATION cmake/Utils/) + install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/cmake/Utils/ DESTINATION cmake/Utils) install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/CMakeLists.txt DESTINATION .) - install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/cmake/tvm_ffi-config.cmake DESTINATION lib/cmake/tvm_ffi/) + install(FILES ${CMAKE_CURRENT_SOURCE_DIR}/cmake/tvm_ffi-config.cmake DESTINATION cmake) endif() ########## Install the related for normal cmake library ########## diff --git a/ffi/cmake/tvm_ffi-config.cmake b/ffi/cmake/tvm_ffi-config.cmake index 003d6dd1e304..01f60ca10bff 100644 --- a/ffi/cmake/tvm_ffi-config.cmake +++ b/ffi/cmake/tvm_ffi-config.cmake @@ -54,3 +54,5 @@ set_target_properties( tvm_ffi_shared PROPERTIES INTERFACE_INCLUDE_DIRECTORIES "${tvm_ffi_INCLUDE_DIR};${tvm_ffi_DLPACK_INCLUDE_DIR}" ) +# extra cmake functions +include(${CMAKE_CURRENT_LIST_DIR}/Utils/Library.cmake) diff --git a/ffi/examples/get_started/README.md b/ffi/examples/get_started/README.md index 746d24ae91f7..002d4375a6dc 100644 --- a/ffi/examples/get_started/README.md +++ b/ffi/examples/get_started/README.md @@ -23,11 +23,10 @@ that can be loaded in different environments. The example implements a simple "add one" operation that adds 1 to each element of an input tensor, showing how to create C++ functions callable from Python. - You can run this quick start example by: ```bash -# ensure you installed tvm-ffi first once +# ensure you installed tvm-ffi first pip install -e ../.. # Build and run the complete example @@ -49,8 +48,8 @@ in Python and C++. ## Compile without CMake -You can also compile the modules directly using using -flags provided by the `tvm-ffi-config` tool +You can also compile the modules directly using +flags provided by the `tvm-ffi-config` tool. ```bash g++ -shared -fPIC `tvm-ffi-config --cxxflags` \ diff --git a/ffi/examples/packaging/CMakeLists.txt b/ffi/examples/packaging/CMakeLists.txt new file mode 100644 index 000000000000..47e5040a0d73 --- /dev/null +++ b/ffi/examples/packaging/CMakeLists.txt @@ -0,0 +1,73 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +cmake_minimum_required(VERSION 3.18) +project(tvm_ffi_extension) + +option(TVM_FFI_EXT_FROM_SOURCE "Build tvm_ffi from source, useful for cross compilation." ON) +option(TVM_FFI_EXT_SHIP_DEBUG_SYMBOLS "Ship debug symbols" ON) + +# There are two ways to include tvm_ffi +# +# 1. Build tvm_ffi from source, which is reasonably cheap since tvm ffi is small +# 2. Use the pre-built tvm_ffi shipped from the pip +# +# This example shows both options, you only need to pick a specific one. +# +# - For common build cases, using pre-built and link tvm_ffi_shared is sufficient. +# - For cases where you may want to cross-compile or bundle part of tvm_ffi_objects directly +# into your project, opt for building tvm_ffi from source path. +# Note that it is always safe to build from source and extra cost of building tvm_ffi is small. +# So when in doubt, you can always choose to the building tvm_ffi from source route. +# +# In python or other cases when we dynamically load libtvm_ffi_shared. Even when you build +# from source, you do not need to ship libtvm_ffi_shared.so built here as they are only +# used to supply the linking information. +# first find python related components +find_package(Python COMPONENTS Interpreter REQUIRED) +if (TVM_FFI_BUILD_FROM_SOURCE) + execute_process( + COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --sourcedir + OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT) + message(STATUS "Building tvm_ffi from source: ${tvm_ffi_ROOT}") + add_subdirectory(${tvm_ffi_ROOT} tvm_ffi) +else() + # call tvm_ffi.config to get the cmake directory and set it to tvm_ffi_ROOT + execute_process( + COMMAND "${Python_EXECUTABLE}" -m tvm_ffi.config --cmakedir + OUTPUT_STRIP_TRAILING_WHITESPACE OUTPUT_VARIABLE tvm_ffi_ROOT) + find_package(tvm_ffi CONFIG REQUIRED) +endif() + +# use the projects as usual +add_library(tvm_ffi_extension SHARED src/extension.cc) +target_link_libraries(tvm_ffi_extension tvm_ffi_header) +target_link_libraries(tvm_ffi_extension tvm_ffi_shared) + +# show as tvm_ffi_extension.so +set_target_properties( + tvm_ffi_extension PROPERTIES PREFIX "" +) + +if (TVM_FFI_EXT_SHIP_DEBUG_SYMBOLS) + # ship debugging symbols for backtrace on macos + tvm_ffi_add_prefix_map(tvm_ffi_extension ${CMAKE_CURRENT_SOURCE_DIR}) + tvm_ffi_add_apple_dsymutil(tvm_ffi_extension) + install(DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/ DESTINATION . FILES_MATCHING PATTERN "*.dSYM") +endif() + +install(TARGETS tvm_ffi_extension DESTINATION .) diff --git a/ffi/examples/packaging/README.md b/ffi/examples/packaging/README.md new file mode 100644 index 000000000000..9535581af622 --- /dev/null +++ b/ffi/examples/packaging/README.md @@ -0,0 +1,61 @@ + + + + + + + + + + + + + + + + + +# TVM FFI Packaging Example + +This is an example project that packages a tvm-ffi based library +into a Python ABI-agnostic wheel. + +This example can also serve as a guideline for general +packaging as well. + +- Source-level build for cross-compilation support in CMake +- Registration via global function table + +## Install the wheel + +```bash +pip install . +``` + +### Note on build and auditwheel + +Note: When running the auditwheel process, make sure to skip +`libtvm_ffi_shared.so` as they are shipped via the tvm_ffi package. + +## Run the example + +After installing the `tvm_ffi_extension` example package, you can run the following example +that invokes the `add_one` function exposed. + +```bash +python run_example.py add_one +``` + +You can also run the following command to see how error is raised and propagated +across the language boundaries. + +```python +python run_example.py raise_error +``` + +When possible, tvm_ffi will try to preserve traceback across language boundary. You will see traceback like +``` +File "src/extension.cc", line 45, in void tvm_ffi_extension::RaiseError(tvm::ffi::String) +``` +If you are in an IDE like VSCode, you can click and jump to the C++ lines of error when +the debug symbols are preserved. diff --git a/ffi/examples/packaging/pyproject.toml b/ffi/examples/packaging/pyproject.toml new file mode 100644 index 000000000000..e38ebeccff4d --- /dev/null +++ b/ffi/examples/packaging/pyproject.toml @@ -0,0 +1,58 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[project] +name = "tvm-ffi-extension" +version = "0.1.0" + +readme = "README.md" +license = { text = "Apache 2.0" } +classifiers = [ + "License :: OSI Approved :: Apache Software License", + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", +] +keywords = ["machine learning", "inference"] +requires-python = ">=3.9" + +dependencies = ["apache-tvm-ffi"] + +[build-system] +requires = ["scikit-build-core>=0.10.0", "apache-tvm-ffi"] +build-backend = "scikit_build_core.build" + +[tool.scikit-build] +# the wheel is abi agnostic +wheel.py-api = "py3" +minimum-version = "build-system.requires" + +# Build configuration +build-dir = "build" +build.verbose = true + +# CMake configuration +cmake.version = "CMakeLists.txt" +cmake.build-type = "RelWithDebugInfo" + +# Logging +logging.level = "INFO" + +# Wheel configuration +wheel.packages = ["python/tvm_ffi_extension"] +wheel.install-dir = "tvm_ffi_extension" diff --git a/ffi/examples/packaging/python/tvm_ffi_extension/__init__.py b/ffi/examples/packaging/python/tvm_ffi_extension/__init__.py new file mode 100644 index 000000000000..4cd4207df136 --- /dev/null +++ b/ffi/examples/packaging/python/tvm_ffi_extension/__init__.py @@ -0,0 +1,48 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations. +from .base import _LIB +from . import _ffi_api + + +def add_one(x, y): + """ + Adds one to the input tensor. + + Parameters + ---------- + x : Tensor + The input tensor. + y : Tensor + The output tensor. + """ + return _LIB.add_one(x, y) + + +def raise_error(msg): + """ + Raises an error with the given message. + + Parameters + ---------- + msg : str + The message to raise the error with. + + Raises + ------ + RuntimeError + The error raised by the function. + """ + return _ffi_api.raise_error(msg) diff --git a/ffi/examples/packaging/python/tvm_ffi_extension/_ffi_api.py b/ffi/examples/packaging/python/tvm_ffi_extension/_ffi_api.py new file mode 100644 index 000000000000..1ab9abd765a8 --- /dev/null +++ b/ffi/examples/packaging/python/tvm_ffi_extension/_ffi_api.py @@ -0,0 +1,24 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations. + +import tvm_ffi + +# make sure lib is loaded first +from .base import _LIB + +# this is a short cut to register all the global functions +# prefixed by `tvm_ffi_extension.` to this module +tvm_ffi._init_api("tvm_ffi_extension", __name__) diff --git a/ffi/examples/packaging/python/tvm_ffi_extension/base.py b/ffi/examples/packaging/python/tvm_ffi_extension/base.py new file mode 100644 index 000000000000..ed73193770a8 --- /dev/null +++ b/ffi/examples/packaging/python/tvm_ffi_extension/base.py @@ -0,0 +1,37 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations. +# Base logic to load library for extension package +import tvm_ffi +import os +import sys + + +def _load_lib(): + # first look at the directory of the current file + file_dir = os.path.dirname(os.path.realpath(__file__)) + + if sys.platform.startswith("win32"): + lib_dll_name = "tvm_ffi_extension.dll" + elif sys.platform.startswith("darwin"): + lib_dll_name = "tvm_ffi_extension.dylib" + else: + lib_dll_name = "tvm_ffi_extension.so" + + lib_path = os.path.join(file_dir, lib_dll_name) + return tvm_ffi.load_module(lib_path) + + +_LIB = _load_lib() diff --git a/ffi/pyproject.toml b/ffi/pyproject.toml index 3efa1d9455a1..8ed9e275e2b3 100644 --- a/ffi/pyproject.toml +++ b/ffi/pyproject.toml @@ -17,7 +17,7 @@ [project] name = "apache-tvm-ffi" -version = "0.1.0a3" +version = "0.1.0a5" description = "tvm ffi" authors = [{ name = "TVM FFI team" }] diff --git a/ffi/python/tvm_ffi/cython/function.pxi b/ffi/python/tvm_ffi/cython/function.pxi index dcd300c9b036..00a0bb351508 100644 --- a/ffi/python/tvm_ffi/cython/function.pxi +++ b/ffi/python/tvm_ffi/cython/function.pxi @@ -56,10 +56,7 @@ def load_torch_get_current_cuda_stream(): return fallback_get_current_cuda_stream -if torch is not None: - # when torch is available, jit compile the get_current_cuda_stream function - # the torch caches the extension so second loading is faster - torch_get_current_cuda_stream = load_torch_get_current_cuda_stream() +torch_get_current_cuda_stream = None cdef inline object make_ret_small_str(TVMFFIAny result): @@ -149,6 +146,8 @@ cdef inline int make_args(tuple py_args, TVMFFIAny* out, list temp_args, if is_cuda and ctx_dev_type != NULL and ctx_dev_type[0] == -1: ctx_dev_type[0] = temp_dltensor.device.device_type ctx_dev_id[0] = temp_dltensor.device.device_id + if torch_get_current_cuda_stream is None: + torch_get_current_cuda_stream = load_torch_get_current_cuda_stream() temp_ptr = torch_get_current_cuda_stream(temp_dltensor.device.device_id) ctx_stream[0] = temp_ptr temp_args.append(arg) diff --git a/ffi/python/tvm_ffi/libinfo.py b/ffi/python/tvm_ffi/libinfo.py index 8974574fe9dd..b449bc1abcf5 100644 --- a/ffi/python/tvm_ffi/libinfo.py +++ b/ffi/python/tvm_ffi/libinfo.py @@ -95,7 +95,7 @@ def find_source_path(): def find_cmake_path(): """Find the preferred cmake path.""" candidates = [ - os.path.join(os.path.dirname(os.path.realpath(__file__)), "lib", "cmake"), + os.path.join(os.path.dirname(os.path.realpath(__file__)), "cmake"), os.path.join(os.path.dirname(os.path.realpath(__file__)), "..", "..", "cmake"), ] for candidate in candidates: