Maintain separate extension_cpp and extension_cpp_stable

Author: mikaylagawarecki

.github/scripts/setup-env.sh

@@ -93,15 +93,23 @@ echo '::group::Install third party dependencies prior to extension-cpp install'
 # - It happily pulls in pre-releases, which can lead to more problems down the line.
 #   `pip` does not unless explicitly told to do so.
 # Thus, we use `easy_install` to extract the third-party dependencies here and install them upfront with `pip`.
+pushd extension_cpp
 python setup.py egg_info
 # The requires.txt cannot be used with `pip install -r` directly. The requirements are listed at the top and the
 # optional dependencies come in non-standard syntax after a blank line. Thus, we just extract the header.
 sed -e '/^$/,$d' *.egg-info/requires.txt | tee requirements.txt
 pip install --progress-bar=off -r requirements.txt
 echo '::endgroup::'
 
-echo '::group::Install extension-cpp'
-python setup.py develop
+echo '::group::Install extension_cpp (standard ATen API)'
+pip install -e . --no-build-isolation
+popd
+echo '::endgroup::'
+
+echo '::group::Install extension_cpp_stable (stable ABI)'
+pushd extension_cpp_stable
+pip install -e . --no-build-isolation
+popd
 echo '::endgroup::'
 
 echo '::group::Collect environment information'

.github/workflows/tests.yml

@@ -19,7 +19,7 @@ jobs:
           - python-version: 3.13
             runner: linux.g5.4xlarge.nvidia.gpu
             gpu-arch-type: cuda
-            gpu-arch-version: "12.4"
+            gpu-arch-version: "12.9"
       fail-fast: false
     uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main
     permissions:

README.md

@@ -1,18 +1,42 @@
 # C++/CUDA Extensions in PyTorch
 
-An example of writing a C++/CUDA extension for PyTorch. See
-[here](https://pytorch.org/tutorials/advanced/cpp_custom_ops.html) for the accompanying tutorial.
-This repo demonstrates how to write an example `extension_cpp.ops.mymuladd`
-custom op that has both custom CPU and CUDA kernels.
+This repository contains two example C++/CUDA extensions for PyTorch:
 
-The examples in this repo work with PyTorch 2.4+.
+1. **extension_cpp** - Uses the standard ATen/LibTorch API
+2. **extension_cpp_stable** - Uses the [LibTorch Stable ABI](https://pytorch.org/docs/main/notes/libtorch_stable_abi.html)
 
-To build:
+Both extensions demonstrate how to write an example `mymuladd` custom op that has both
+custom CPU and CUDA kernels.
+
+## extension_cpp (Standard ATen API)
+
+Uses the full ATen/LibTorch API. This is the traditional way of writing PyTorch extensions.
+See [this tutorial](https://pytorch.org/tutorials/advanced/cpp_custom_ops.html) for more details.
+
+## extension_cpp_stable (Stable ABI)
+
+Uses the LibTorch Stable ABI to ensure that the extension built can be run with any version
+of PyTorch >= 2.10.0, without needing to recompile for each PyTorch version.
+
+The `extension_cpp_stable` examples require PyTorch 2.10+.
+
+## Building
+
+To build extension_cpp (standard API):
+```
+cd extension_cpp
+pip install --no-build-isolation -e .
 ```
+
+To build extension_cpp_stable (stable ABI):
+```
+cd extension_cpp_stable
 pip install --no-build-isolation -e .
 ```
 
-To test:
+## Testing
+
+To test both extensions:
 ```
 python test/test_extension.py
 ```

extension_cpp/__init__.py extension_cpp/extension_cpp/__init__.pyextension_cpp/__init__.py renamed to extension_cpp/extension_cpp/__init__.py

@@ -79,7 +79,9 @@ def get_extensions():
     ext_modules=get_extensions(),
     install_requires=["torch"],
     description="Example of PyTorch C++ and CUDA extensions",
-    long_description=open("README.md").read(),
+    long_description=open(
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "README.md")
+    ).read(),
     long_description_content_type="text/markdown",
     url="https://github.com/pytorch/extension-cpp",
     cmdclass={"build_ext": BuildExtension},