[FFI][DOCS] Wheel Packaging

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 ##########

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)

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`  \

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

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.

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"

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)

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

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

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" }]

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)

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: