This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new f703a0c feat: Build PyTorch DLPack extension ahead of time (#189)
f703a0c is described below
commit f703a0cf9358fa30d8faee719f905c58d8ca6ee3
Author: Zihao Ye <[email protected]>
AuthorDate: Fri Oct 31 16:11:39 2025 -0700
feat: Build PyTorch DLPack extension ahead of time (#189)
This PR adds option to build PyTorch DLPack extension ahead of time.
Backward compatible: Default behavior unchanged, opt-in only.
Remaining work includes adding the release workflow and @cyx-6 could
take over, you can refer to
https://github.com/flashinfer-ai/flashinfer/blob/main/.github/workflows/release.yml
---------
Signed-off-by: Yaoyao Ding <[email protected]>
Co-authored-by: Yaxing Cai <[email protected]>
Co-authored-by: Yaoyao Ding <[email protected]>
---
addons/torch_c_dlpack_ext/build_backend.py | 87 ++++++++++++++++++++++
addons/torch_c_dlpack_ext/pyproject.toml | 38 ++++++++++
.../torch_c_dlpack_ext/__init__.py | 19 +++++
.../torch_c_dlpack_ext/torch_c_dlpack_ext/core.py | 53 +++++++++++++
python/tvm_ffi/_optional_torch_c_dlpack.py | 7 ++
5 files changed, 204 insertions(+)
diff --git a/addons/torch_c_dlpack_ext/build_backend.py
b/addons/torch_c_dlpack_ext/build_backend.py
new file mode 100644
index 0000000..8a14388
--- /dev/null
+++ b/addons/torch_c_dlpack_ext/build_backend.py
@@ -0,0 +1,87 @@
+# 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.
+"""build backend for torch c dlpack ext."""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+from pathlib import Path
+
+from setuptools import build_meta as orig
+
+_root = Path(__file__).parent.resolve()
+_package_path = _root / "torch_c_dlpack_ext"
+
+
+get_requires_for_build_sdist = orig.get_requires_for_build_sdist
+get_requires_for_build_editable = orig.get_requires_for_build_editable
+prepare_metadata_for_build_wheel = orig.prepare_metadata_for_build_wheel
+prepare_metadata_for_build_editable = orig.prepare_metadata_for_build_editable
+build_sdist = orig.build_sdist
+build_editable = orig.build_editable
+
+
+def _is_lib_prebuilt() -> bool:
+ if sys.platform.startswith("win32"):
+ extension = "dll"
+ elif sys.platform.startswith("darwin"):
+ extension = "dylib"
+ else:
+ extension = "so"
+ return next(_package_path.rglob(f"*.{extension}"), None) is not None
+
+
+def get_requires_for_build_wheel(
+ config_settings: orig._ConfigSettings = None,
+) -> list[str]:
+ """Get build requirements for wheel, conditionally including torch and
apache-tvm-ffi."""
+ requires = orig.get_requires_for_build_wheel(config_settings)
+ if not _is_lib_prebuilt():
+ # build wheel from sdist package, requires apache-tvm-ffi>=0.1.1 to
build lib
+ requires.append("apache-tvm-ffi>=0.1.1")
+ return requires
+
+
+def build_wheel(
+ wheel_directory: orig.StrPath,
+ config_settings: orig._ConfigSettings = None,
+ metadata_directory: orig.StrPath | None = None,
+) -> str:
+ """Build wheel."""
+ if not _is_lib_prebuilt():
+ # build wheel from sdist package, compile the torch c dlpack ext
library locally.
+ import torch # noqa: PLC0415
+
+ if hasattr(torch.Tensor, "__c_dlpack_exchange_api__"):
+ print(
+ "torch.Tensor already has attribute __c_dlpack_exchange_api__.
"
+ "No need to build any torch c dlpackc libs."
+ )
+ else:
+ subprocess.run(
+ [
+ sys.executable,
+ "-m",
+ "tvm_ffi.utils._build_optional_torch_c_dlpack",
+ "--output-dir",
+ str(_package_path),
+ "--build-with-cuda" if torch.cuda.is_available() else "",
+ ],
+ check=True,
+ )
+ return orig.build_wheel(wheel_directory, config_settings,
metadata_directory)
diff --git a/addons/torch_c_dlpack_ext/pyproject.toml
b/addons/torch_c_dlpack_ext/pyproject.toml
new file mode 100644
index 0000000..549dc85
--- /dev/null
+++ b/addons/torch_c_dlpack_ext/pyproject.toml
@@ -0,0 +1,38 @@
+# 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 = "torch_c_dlpack_ext"
+version = "0.1.0"
+requires-python = ">=3.9"
+description = "torch c dlpack ext"
+dependencies = ["torch"]
+
+[build-system]
+requires = ["setuptools>=61.0", "torch"]
+build-backend = "build_backend"
+backend-path = ["."]
+
+[tool.setuptools]
+include-package-data = true
+py-modules = ["build_backend"]
+
+[tool.setuptools.packages.find]
+include = ["torch_c_dlpack_ext*"]
+
+[tool.setuptools.package-data]
+torch_c_dlpack_ext = ["*.so", "*.dll", "*.dylib"]
diff --git a/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/__init__.py
b/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/__init__.py
new file mode 100644
index 0000000..cbde1de
--- /dev/null
+++ b/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/__init__.py
@@ -0,0 +1,19 @@
+# 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.
+"""torch c dlpack ext pakcage."""
+
+import core
diff --git a/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/core.py
b/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/core.py
new file mode 100644
index 0000000..a2030ea
--- /dev/null
+++ b/addons/torch_c_dlpack_ext/torch_c_dlpack_ext/core.py
@@ -0,0 +1,53 @@
+# 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.
+"""torch c dlpack ext core methods."""
+
+import ctypes
+import sys
+from pathlib import Path
+
+import torch
+from packaging.version import Version
+
+
+def load_torch_c_dlpack_extension() -> None:
+ """Load the torch c dlpack extension based on torch version."""
+ if hasattr(torch.Tensor, "__c_dlpack_exchange_api__"):
+ return None
+ version = Version(torch.__version__)
+ if sys.platform.startswith("win32"):
+ extension = "dll"
+ elif sys.platform.startswith("darwin"):
+ extension = "dylib"
+ else:
+ extension = "so"
+ suffix = "cuda" if torch.cuda.is_available() else "cpu"
+ lib_path = (
+ Path(__file__).parent
+ /
f"libtorch_c_dlpack_addon_torch{version.major}{version.minor}-{suffix}.{extension}"
+ )
+ if not lib_path.exists() or not lib_path.is_file():
+ raise ImportError("No matching prebuilt torch c dlpack extension")
+ lib = ctypes.CDLL(str(lib_path))
+ func = lib.TorchDLPackExchangeAPIPtr
+ func.restype = ctypes.c_uint64
+ func.argtypes = []
+ setattr(torch.Tensor, "__c_dlpack_exchange_api__", func())
+ return lib
+
+
+_lib = load_torch_c_dlpack_extension()
diff --git a/python/tvm_ffi/_optional_torch_c_dlpack.py
b/python/tvm_ffi/_optional_torch_c_dlpack.py
index 2be2301..4c5e024 100644
--- a/python/tvm_ffi/_optional_torch_c_dlpack.py
+++ b/python/tvm_ffi/_optional_torch_c_dlpack.py
@@ -54,6 +54,13 @@ def load_torch_c_dlpack_extension() -> Any:
return None
"""Load the torch c dlpack extension."""
+ try:
+ import torch_c_dlpack_ext # type: ignore # noqa: PLC0415, F401
+
+ return None
+ except ImportError:
+ pass
+
try:
# todo: check whether a prebuilt package is installed, if so, use it.
...
