mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
327e07ac1d
This PR adds two main parts:
- shim.h stable C APIs into torch::Library APIs
- a higher level API in torch/csrc/stable/library.h that calls into this shim.h + otherwise is self contained
Goal: custom kernel writers should be able to call the apis in the directories above in order to register their library in a way that allows their custom extension to run with a different libtorch version than it was built with.
Subplots resolved:
- Do we want a whole separate StableLibrary or do we want to freeze torch::Library and add `m.stable_impl(cstring, void (*fn)(void **, int64_t, int64_t)` into it
- Yes, we want a separate StableLibrary. We cannot freeze Library and it is NOT header only.
- Should I use unint64_t as the common denominator instead of void* to support 32bit architectures better?
- Yes, and done
- Should I add a stable `def` and `fragment` when those can be done in python?
- I think we do want these --- and now they're done
- Where should library_stable_impl.cpp live? -- no longer relevant
- I need some solid test cases to make sure everything's going ok. I've intentionally thrown in a bunch of random dtypes into the signature, but I still haven't tested returning multiple things, returning nothing, complex dtypes, etc.
- Have since tested all the torch library endpoints. the others can be tested in a followup to separate components that need to be in shim.h vs can be added later
Pull Request resolved: https://github.com/pytorch/pytorch/pull/148124
Approved by: https://github.com/albanD, https://github.com/zou3519
68 lines
1.8 KiB
Python
68 lines
1.8 KiB
Python
import distutils.command.clean
|
|
import shutil
|
|
from pathlib import Path
|
|
|
|
from setuptools import find_packages, setup
|
|
|
|
from torch.utils.cpp_extension import BuildExtension, CppExtension
|
|
|
|
|
|
ROOT_DIR = Path(__file__).parent
|
|
CSRC_DIR = ROOT_DIR / "libtorch_agnostic" / "csrc"
|
|
|
|
|
|
class clean(distutils.command.clean.clean):
|
|
def run(self):
|
|
# Run default behavior first
|
|
distutils.command.clean.clean.run(self)
|
|
|
|
# Remove extension
|
|
for path in (ROOT_DIR / "libtorch_agnostic").glob("**/*.so"):
|
|
path.unlink()
|
|
# Remove build and dist and egg-info directories
|
|
dirs = [
|
|
ROOT_DIR / "build",
|
|
ROOT_DIR / "dist",
|
|
ROOT_DIR / "libtorch_agnostic.egg-info",
|
|
]
|
|
for path in dirs:
|
|
if path.exists():
|
|
shutil.rmtree(str(path), ignore_errors=True)
|
|
|
|
|
|
def get_extension():
|
|
extra_compile_args = {
|
|
"cxx": ["-fdiagnostics-color=always"],
|
|
}
|
|
|
|
sources = list(CSRC_DIR.glob("**/*.cpp"))
|
|
|
|
return [
|
|
CppExtension(
|
|
"libtorch_agnostic._C",
|
|
sources=sorted(str(s) for s in sources),
|
|
py_limited_api=True,
|
|
extra_compile_args=extra_compile_args,
|
|
extra_link_args=[],
|
|
)
|
|
]
|
|
|
|
|
|
setup(
|
|
name="libtorch_agnostic",
|
|
version="0.0",
|
|
author="PyTorch Core Team",
|
|
description="Example of libtorch agnostic extension",
|
|
packages=find_packages(exclude=("test",)),
|
|
package_data={"libtorch_agnostic": ["*.dll", "*.dylib", "*.so"]},
|
|
install_requires=[
|
|
"torch",
|
|
],
|
|
ext_modules=get_extension(),
|
|
cmdclass={
|
|
"build_ext": BuildExtension.with_options(no_python_abi_suffix=True),
|
|
"clean": clean,
|
|
},
|
|
options={"bdist_wheel": {"py_limited_api": "cp39"}},
|
|
)
|