This is an automated email from the ASF dual-hosted git repository.
yongwww pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 4ec17095d5 [FFI][DOCS] Wheel Packaging (#18256)
4ec17095d5 is described below
commit 4ec17095d5498546ac7efeaa805b8860b613e10c
Author: Tianqi Chen <[email protected]>
AuthorDate: Sat Aug 30 23:12:45 2025 -0400
[FFI][DOCS] Wheel Packaging (#18256)
[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 +++++++++++++++++
.../packaging/python/tvm_ffi_extension/__init__.py | 48 ++++++++++++++
.../packaging/python/tvm_ffi_extension/_ffi_api.py | 24 +++++++
.../packaging/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(-)
diff --git a/ffi/CMakeLists.txt b/ffi/CMakeLists.txt
index f40313636a..90f1f89cbb 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 003d6dd1e3..01f60ca10b 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 746d24ae91..002d4375a6 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 0000000000..47e5040a0d
--- /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 0000000000..9535581af6
--- /dev/null
+++ b/ffi/examples/packaging/README.md
@@ -0,0 +1,61 @@
+<!--- 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. -->
+
+# 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 0000000000..e38ebeccff
--- /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 0000000000..4cd4207df1
--- /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 0000000000..1ab9abd765
--- /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 0000000000..ed73193770
--- /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 3efa1d9455..8ed9e275e2 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 dcd300c9b0..00a0bb3515 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] = <TVMFFIStreamHandle>temp_ptr
temp_args.append(arg)
diff --git a/ffi/python/tvm_ffi/libinfo.py b/ffi/python/tvm_ffi/libinfo.py
index 8974574fe9..b449bc1abc 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:
