diff --git a/.github/workflows/build_and_deploy.yml b/.github/workflows/build_and_deploy.yml new file mode 100644 index 00000000..e5ab657f --- /dev/null +++ b/.github/workflows/build_and_deploy.yml @@ -0,0 +1,195 @@ +name: Build and Deploy Manylinux Wheels + +# Trigger workflow on PRs, tags, and manual dispatch +on: + push: + tags: + - 'v*' + pull_request: + branches: + - main + workflow_dispatch: + inputs: + torch_version: + description: 'PyTorch version to build (e.g., 2.4.0)' + required: false + default: '' + cuda_version: + description: 'CUDA version to build (e.g., cu121, cpu)' + required: false + default: '' + +env: + CUDA_ARCH_MAP: '{"11.8": "8.0;8.6;8.9;+PTX", "12.1": "8.0;8.6;8.9;9.0;+PTX", "12.4": "8.0;8.6;8.9;9.0;+PTX", "12.6": "8.0;8.6;8.9;9.0;+PTX", "12.8": "8.0;8.6;8.9;9.0;10.0;12.0;+PTX", "none": ""}' + +jobs: + build-manylinux-wheels: + runs-on: ubuntu-latest + strategy: + matrix: + include: + # # PyTorch 2.4.x with multiple CUDA versions + # - torch-version: '2.4.0' + # torch-cuda: 'cpu' + # cuda-version: 'none' + # - torch-version: '2.4.0' + # torch-cuda: 'cu118' + # cuda-version: '11.8' + # - torch-version: '2.4.0' + # torch-cuda: 'cu121' + # cuda-version: '12.1' + # - torch-version: '2.4.0' + # torch-cuda: 'cu124' + # cuda-version: '12.4' + + # # PyTorch 2.5.x with multiple CUDA versions + # - torch-version: '2.5.0' + # torch-cuda: 'cpu' + # cuda-version: 'none' + # - torch-version: '2.5.0' + # torch-cuda: 'cu118' + # cuda-version: '11.8' + # - torch-version: '2.5.0' + # torch-cuda: 'cu121' + # cuda-version: '12.1' + # - torch-version: '2.5.0' + # torch-cuda: 'cu124' + # cuda-version: '12.4' + + # PyTorch 2.6.x with multiple CUDA versions + - torch-version: '2.6.0' + torch-cuda: 'cpu' + cuda-version: 'none' + # - torch-version: '2.6.0' + # torch-cuda: 'cu118' + # cuda-version: '11.8' + # - torch-version: '2.6.0' + # torch-cuda: 'cu124' + # cuda-version: '12.4' + # - torch-version: '2.6.0' + # torch-cuda: 'cu126' + # cuda-version: '12.6' + + # PyTorch 2.7.x with multiple CUDA versions + - torch-version: '2.7.0' + torch-cuda: 'cpu' + cuda-version: 'none' + # - torch-version: '2.7.0' + # torch-cuda: 'cu118' + # cuda-version: '11.8' + # - torch-version: '2.7.0' + # torch-cuda: 'cu126' + # cuda-version: '12.6' + # - torch-version: '2.7.0' + # torch-cuda: 'cu128' + # cuda-version: '12.8' + + steps: + - uses: actions/checkout@v4 + + # - name: Free Disk Space + # run: | + # sudo rm -rf /usr/local/.ghcup + # sudo rm -rf /opt/hostedtoolcache/CodeQL + # sudo rm -rf /usr/local/lib/android/sdk/ndk + # sudo rm -rf /usr/share/dotnet + # sudo rm -rf /opt/ghc + # sudo rm -rf /usr/local/share/boost + # sudo apt-get clean + # echo "Disk space after cleanup:" + # df -h + + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.11' + + - name: Install cibuildwheel + run: | + python -m pip install cibuildwheel + + - name: Set CUDA architecture list + run: | + ARCH_LIST=$(echo "$CUDA_ARCH_MAP" | jq -r '.["${{ matrix.cuda-version }}"]') + echo "TORCH_CUDA_ARCH_LIST=$ARCH_LIST" >> $GITHUB_ENV + + - name: Extract base version from pyproject.toml + run: | + BASE_VERSION=$(python -c "import tomllib; print(tomllib.load(open('pyproject.toml', 'rb'))['tool']['setuptools_scm']['fallback_version'])" 2>/dev/null || \ + python -c "import tomli; print(tomli.load(open('pyproject.toml', 'rb'))['tool']['setuptools_scm']['fallback_version'])") + echo "BASE_VERSION=$BASE_VERSION" >> $GITHUB_ENV + echo "Base version: $BASE_VERSION" + + - name: Build manylinux wheels + env: + CIBW_BEFORE_BUILD: | + pip install --upgrade pip setuptools wheel setuptools-scm + if [ "${{ matrix.torch-cuda }}" = "cpu" ]; then + pip install --no-cache-dir torch==${{ matrix.torch-version }}+cpu --index-url https://download.pytorch.org/whl/cpu + else + pip install --no-cache-dir torch==${{ matrix.torch-version }}+${{ matrix.torch-cuda }} --index-url https://download.pytorch.org/whl/${{ matrix.torch-cuda }} + fi + pip install numpy + export SETUPTOOLS_SCM_PRETEND_VERSION="${{ env.BASE_VERSION }}+torch${{ matrix.torch-version }}.${{ matrix.torch-cuda }}" + export TORCH_HARMONICS_ENABLE_OPENMP=1 + if [ "${{ matrix.cuda-version }}" != "none" ]; then + export CUDA_HOME=/usr/local/cuda-${{ matrix.cuda-version }} + export PATH=$CUDA_HOME/bin:$PATH + export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH + fi + CIBW_BEFORE_TEST: | + echo "Installing PyTorch for testing..." + if [ "${{ matrix.torch-cuda }}" = "cpu" ]; then + pip install torch==${{ matrix.torch-version }}+cpu --index-url https://download.pytorch.org/whl/cpu + else + pip install torch==${{ matrix.torch-version }}+${{ matrix.torch-cuda }} --index-url https://download.pytorch.org/whl/${{ matrix.torch-cuda }} + fi + pip install numpy + if [ "${{ matrix.cuda-version }}" != "none" ]; then + export CUDA_HOME=/usr/local/cuda-${{ matrix.cuda-version }} + export PATH=$CUDA_HOME/bin:$PATH + export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH + fi + CIBW_TEST_INSIDE: true + TORCH_CUDA_VERSION: ${{ matrix.cuda-version }} + TORCH_CUDA_ARCH_LIST: ${{ env.TORCH_CUDA_ARCH_LIST }} + CIBW_ENVIRONMENT: "SETUPTOOLS_SCM_PRETEND_VERSION=${{ env.BASE_VERSION }}+torch${{ matrix.torch-version }}.${{ matrix.torch-cuda }} TORCH_CUDA_VERSION=${{ matrix.cuda-version }}" + + run: | + cibuildwheel --platform linux + + - name: Inspect built wheels and .so files + if: always() + run: | + echo "Listing wheelhouse contents:" + ls -lh wheelhouse/ + echo "Checking .so symbols and ABI:" + for so in wheelhouse/*.so; do + echo "==== $so ====" + nm -D "$so" | grep exception_ptr || echo "No exception_ptr symbols found" + strings "$so" | grep ABI || echo "No ABI macro found in $so" + done + + - name: Upload wheel artifacts + uses: actions/upload-artifact@v4 + with: + name: wheel-torch-${{ matrix.torch-version }}-${{ matrix.torch-cuda }} + path: wheelhouse/*.whl + + - name: Upload to TestPyPI (for testing) + if: github.event_name == 'workflow_dispatch' || github.event_name == 'pull_request' + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.TEST_PYPI_PASSWORD }} + run: | + python -m pip install twine + python -m twine upload --verbose --repository testpypi wheelhouse/*.whl + + - name: Upload to PyPI (production) + if: startsWith(github.ref, 'refs/tags/') + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} + run: | + python -m pip install twine + python -m twine upload wheelhouse/*.whl \ No newline at end of file diff --git a/.github/workflows/deploy_pypi.yml b/.github/workflows/deploy_pypi.yml deleted file mode 100644 index 6bda954c..00000000 --- a/.github/workflows/deploy_pypi.yml +++ /dev/null @@ -1,39 +0,0 @@ -name: Deploy to Pypi - -on: - release: - types: - - published - workflow_dispatch: - -jobs: - pypi-publish: - name: Publish release to PyPI - runs-on: ubuntu-latest - environment: - name: pypi - url: https://pypi.org/p/torch-harmonics - - steps: - - uses: actions/checkout@v4 - - - name: Set up Python 3.9 - uses: actions/setup-python@v5 - with: - python-version: "3.9" - - - name: Install build dependencies - run: | - python3 -m pip install --upgrade pip setuptools build wheel - python3 -m pip install numpy - python3 -m pip install torch --extra-index-url https://download.pytorch.org/whl/cpu - - - name: Build distribution - run: | - python3 -m build --sdist - - - name: Publish package to PyPI - uses: pypa/gh-action-pypi-publish@release/v1 - with: - user: __token__ - password: ${{ secrets.PYPI_PASSWORD }} \ No newline at end of file diff --git a/.gitignore b/.gitignore index eef5e947..9f2e5626 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,36 @@ *.DS_Store __pycache__ *.so -checkpoints \ No newline at end of file +checkpoints + +# Build artifacts +build/ +dist/ +wheelhouse/ +*.egg-info/ + +# Compiled extensions (keep source, ignore compiled) +*.so +*.pyd + +# CUDA compilation artifacts +*.cubin +*.fatbin +*.ptx + +# Python cache +*.pyc +*.pyo +*.pyd +__pycache__/ +*.egg-info/ + +# IDE +.vscode/ +.idea/ +*.swp +*.swo + +# OS +.DS_Store +Thumbs.db \ No newline at end of file diff --git a/README.md b/README.md index 15ecb3fc..a0a918be 100644 --- a/README.md +++ b/README.md @@ -75,27 +75,36 @@ torch-harmonics has been used to implement a variety of differentiable PDE solve ## Installation -A simple installation can be directly done from PyPI: +In general we recommend building the package yourself with the exact PyTorch build you rely on (see the instructions below) to avoid binary incompatibilities. Many features in torch-harmonics make use of accelerated CUDA kernels which need to match the CUDA packaged with your PyTorch installation. To do so, you may build torch-harmonics using ```bash -pip install torch-harmonics +export FORCE_CUDA_EXTENSION=1 +export TORCH_CUDA_ARCH_LIST="8.0 8.6 8.7 9.0 10.0 12.0+PTX" +pip install --no-binary=torch-harmonics --no-build-isolation torch-harmonics ``` -If you are planning to use spherical convolutions, we recommend building the corresponding custom CUDA kernels. To enforce this, you can set the `FORCE_CUDA_EXTENSION` flag. You may also want to set appropriate architectures with the `TORCH_CUDA_ARCH_LIST` flag. Finally, make sure to disable build isolation via the `--no-build-isolation` flag to ensure that the custom kernels are built with the existing torch installation. +The `--no-build-isolation` flag is necessary to ensure that torch-harmonics extensions are built with the correct PyTorch libraries. For container builds, we recommend setting the flags `FORCE_CUDA_EXTENSION` and `TORCH_CUDA_ARCH_LIST`, as the containers may run on systems with different GPUs than those available on the host system building the container. + +:warning: Please note that the custom CUDA extensions only support CUDA architectures >= 7.0. + +A simple installation using one of the prebuilt wheels on PyPI is also possible. However, it is important that it matches your local PyTorch and CUDA build. To do so, use the helper command below to print the exact `pip install` command for your environment: + ```bash -export FORCE_CUDA_EXTENSION=1 -export TORCH_CUDA_ARCH_LIST="7.0 7.2 7.5 8.0 8.6 8.7 9.0+PTX" -pip install --no-build-isolation torch-harmonics +python -c "import torch, os; parts=torch.__version__.split('+', 1); torch_tag=parts[0]; cuda_tag=parts[1] if len(parts) > 1 else 'cpu'; print(f'pip install torch-harmonics==0.8.1+torch{torch_tag}.{cuda_tag}')" ``` -:warning: Please note that the custom CUDA extensions currently only support CUDA architectures >= 7.0. +This will generate a command of the form +```bash +pip install torch-harmonics==0.8.1+torch{torch_tag}.{cuda_tag} +``` +specifying the exact wheel, matching both PyTorch installation and its CUDA version. If you want to actively develop torch-harmonics, we recommend building it in your environment from github: ```bash git clone git@github.com:NVIDIA/torch-harmonics.git cd torch-harmonics -pip install -e . +pip install --no-build-isolation -e . ``` - + ## More about torch-harmonics diff --git a/build_wheels.py b/build_wheels.py new file mode 100644 index 00000000..0262ef02 --- /dev/null +++ b/build_wheels.py @@ -0,0 +1,180 @@ +#!/usr/bin/env python3 +""" +Build wheels following PyTorch ecosystem naming convention. +This script creates wheels with proper +cuXXX or +cpu suffixes. +""" + +import os +import sys +import subprocess +import shutil +from pathlib import Path + +# PyTorch/CUDA combinations following PyTorch ecosystem conventions +BUILD_MATRIX = [ + # PyTorch 2.4.x with multiple CUDA versions + {"torch": "2.4.0", "cuda": "cu118", "cuda_version": "11.8"}, + {"torch": "2.4.0", "cuda": "cu121", "cuda_version": "12.1"}, + {"torch": "2.4.0", "cuda": "cu124", "cuda_version": "12.4"}, + {"torch": "2.4.0", "cuda": "cpu", "cuda_version": "none"}, + + # PyTorch 2.5.x with multiple CUDA versions + {"torch": "2.5.0", "cuda": "cu118", "cuda_version": "11.8"}, + {"torch": "2.5.0", "cuda": "cu121", "cuda_version": "12.1"}, + {"torch": "2.5.0", "cuda": "cu124", "cuda_version": "12.4"}, + {"torch": "2.5.0", "cuda": "cpu", "cuda_version": "none"}, + + # PyTorch 2.6.x with multiple CUDA versions + {"torch": "2.6.0", "cuda": "cu118", "cuda_version": "11.8"}, + {"torch": "2.6.0", "cuda": "cu121", "cuda_version": "12.1"}, + {"torch": "2.6.0", "cuda": "cu124", "cuda_version": "12.4"}, + {"torch": "2.6.0", "cuda": "cu126", "cuda_version": "12.6"}, + {"torch": "2.6.0", "cuda": "cpu", "cuda_version": "none"}, + +] + +def get_version(): + """Get the current version from pyproject.toml or setup.py.""" + try: + # Try to get version from git tag + result = subprocess.run(['git', 'describe', '--tags', '--dirty'], + capture_output=True, text=True, check=True) + version = result.stdout.strip() + # Remove 'v' prefix if present + if version.startswith('v'): + version = version[1:] + return version + except: + return "0.8.1" # Fallback version + +def build_wheel_for_config(config): + """Build wheel for a specific PyTorch/CUDA configuration.""" + torch_version = config["torch"] + cuda_suffix = config["cuda"] + cuda_version = config["cuda_version"] + + print(f"\n{'='*80}") + print(f"Building wheel for PyTorch {torch_version} + {cuda_suffix}") + print(f"{'='*80}") + + # Create isolated environment + env_name = f"build_env_{torch_version.replace('.', '_')}_{cuda_suffix}" + + try: + # Create virtual environment + subprocess.run([sys.executable, "-m", "venv", env_name], check=True) + + # Determine pip path + if os.name == 'nt': # Windows + pip_cmd = f"{env_name}\\Scripts\\pip" + python_cmd = f"{env_name}\\Scripts\\python" + else: # Unix/Linux/macOS + pip_cmd = f"{env_name}/bin/pip" + python_cmd = f"{env_name}/bin/python" + + # Install build dependencies + subprocess.run([ + pip_cmd, "install", "--upgrade", + "pip", "setuptools", "wheel", "cibuildwheel", "setuptools-scm" + ], check=True) + + # Install PyTorch + if cuda_suffix == "cpu": + subprocess.run([ + pip_cmd, "install", + f"torch=={torch_version}+cpu", + "--index-url", "https://download.pytorch.org/whl/cpu", + "numpy" + ], check=True) + else: + subprocess.run([ + pip_cmd, "install", + f"torch=={torch_version}+{cuda_suffix}", + "--index-url", f"https://download.pytorch.org/whl/{cuda_suffix}", + "numpy" + ], check=True) + + # Build manylinux wheels using cibuildwheel + env = os.environ.copy() + env.update({ + "CIBW_BUILD": "cp39-* cp310-* cp311-* cp312-*", + "CIBW_PLATFORM": "linux", + "CIBW_BEFORE_BUILD": f"pip install torch=={torch_version}+{cuda_suffix} --index-url https://download.pytorch.org/whl/{cuda_suffix} numpy", + "CIBW_ENVIRONMENT": f"TORCH_CUDA_VERSION={cuda_version}", + "CIBW_REPAIR_WHEEL_COMMAND": "auditwheel repair -w {dest_dir} {wheel}", + "CIBW_TEST_COMMAND": "python -c 'import torch_harmonics; print(\"Import successful\")'" + }) + + subprocess.run([ + python_cmd, "-m", "cibuildwheel", "--platform", "linux" + ], check=True, env=env) + + # Rename wheels to follow PyTorch ecosystem convention + version = get_version() + wheel_pattern = f"torch_harmonics-{version}-*.whl" + wheel_files = list(Path("wheelhouse").glob(wheel_pattern)) + + # Rename all wheels in wheelhouse + for wheel_file in wheel_files: + # Extract the wheel filename parts + parts = wheel_file.stem.split('-') + if len(parts) >= 4: + # Format: torch_harmonics-0.8.1-cp310-cp310-linux_x86_64.whl + # Convert to: torch_harmonics-0.8.1+cu121-cp310-cp310-linux_x86_64.whl + new_parts = [ + parts[0], # torch_harmonics + f"{parts[1]}+{cuda_suffix}", # 0.8.1+cu121 + ] + parts[2:] # cp310-cp310-linux_x86_64 + + new_wheel_name = '-'.join(new_parts) + '.whl' + new_wheel_path = Path("wheelhouse") / new_wheel_name + + # Rename the wheel + shutil.move(str(wheel_file), str(new_wheel_path)) + print(f"✅ Renamed wheel to: {new_wheel_name}") + + print(f"✅ Successfully built wheel for PyTorch {torch_version} + {cuda_suffix}") + return True + + except subprocess.CalledProcessError as e: + print(f"❌ Failed to build wheel for PyTorch {torch_version} + {cuda_suffix}: {e}") + return False + finally: + # Clean up environment + if Path(env_name).exists(): + shutil.rmtree(env_name) + +def main(): + """Build wheels for all configurations.""" + print("Building torch-harmonics wheels following PyTorch ecosystem naming convention...") + + successful_builds = [] + failed_builds = [] + + for config in BUILD_MATRIX: + if build_wheel_for_config(config): + successful_builds.append(f"{config['torch']}+{config['cuda']}") + else: + failed_builds.append(f"{config['torch']}+{config['cuda']}") + + # Summary + print(f"\n{'='*80}") + print("BUILD SUMMARY") + print(f"{'='*80}") + print(f"✅ Successful builds: {len(successful_builds)}") + for build in successful_builds: + print(f" - PyTorch {build}") + + if failed_builds: + print(f"❌ Failed builds: {len(failed_builds)}") + for build in failed_builds: + print(f" - PyTorch {build}") + + print(f"\nWheels are available in the 'wheelhouse/' directory") + print("Example wheel names:") + print(" - torch_harmonics-0.8.1+cu121-cp310-cp310-linux_x86_64.whl") + print(" - torch_harmonics-0.8.1+cpu-cp310-cp310-linux_x86_64.whl") + print("\nUpload to PyPI with: python -m twine upload wheelhouse/*") + +if __name__ == "__main__": + main() diff --git a/pyproject.toml b/pyproject.toml index 1189dcd8..c405061c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = [ "setuptools", "setuptools-scm"] +requires = [ "setuptools", "setuptools-scm", "torch"] build-backend = "setuptools.build_meta" [project] @@ -42,8 +42,8 @@ dependencies = [ ] -[tool.setuptools.dynamic] -version = {attr = "torch_harmonics.__version__"} +[tool.setuptools_scm] +fallback_version = "0.8.1" [tool.setuptools.packages.find] include = ["torch_harmonics*"] @@ -62,4 +62,27 @@ dev = [ ] [tool.black] -line-length = 180 \ No newline at end of file +line-length = 180 + +[tool.cibuildwheel] +# Build for Python 3.10, 3.11, 3.12 +build = ["cp39-*", "cp310-*", "cp311-*", "cp312-*"] +build-frontend = "pip; args: --no-build-isolation" + +# Use manylinux for Linux wheels +manylinux-x86_64-image = "manylinux_2_28" +manylinux-i686-image = "manylinux_2_28" + +# Build options + +# Test the built wheel +test-command = [ + "python -c 'import torch; print(torch._C._GLIBCXX_USE_CXX11_ABI)'", + "python -c 'import torch_harmonics; print(\"torch_harmonics imported successfully\"); from torch_harmonics.disco import optimized_kernels_is_available; print(f\"optimized kernels are available: {optimized_kernels_is_available()}\")'" +] + +# Repair wheels with auditwheel - exclude all system, PyTorch, and CUDA libraries +repair-wheel-command = "auditwheel repair --exclude 'libc.so*' --exclude 'libm.so*' --exclude 'libpthread.so*' --exclude 'libdl.so*' --exclude 'libutil.so*' --exclude 'libnsl.so*' --exclude 'libresolv.so*' --exclude 'libcrypt.so*' --exclude 'libnss.so*' --exclude 'libselinux.so*' --exclude 'libpcre.so*' --exclude 'libz.so*' --exclude 'libbz2.so*' --exclude 'liblzma.so*' --exclude 'libtinfo.so*' --exclude 'libncurses.so*' --exclude 'libreadline.so*' --exclude 'libhistory.so*' --exclude 'libform.so*' --exclude 'libmenu.so*' --exclude 'libpanel.so*' --exclude 'libgcc_s.so*' --exclude 'libstdc++.so*' --exclude 'libopenblas.so*' --exclude 'liblapack.so*' --exclude 'libgfortran.so*' --exclude 'libquadmath.so*' --exclude 'libtorch*.so*' --exclude 'libc10.so*' --exclude 'libcudart.so*' --exclude 'libcublas.so*' --exclude 'libcurand.so*' --exclude 'libcusparse.so*' --exclude 'libcufft.so*' --exclude 'libcusolver.so*' --exclude 'libnvrtc.so*' --exclude 'libnvjpeg.so*' --exclude 'libcudnn.so*' -w {dest_dir} {wheel}" + +# Skip building for platforms we don't support +skip = ["*-win32", "*-win_amd64", "*-macosx_arm64", "*-musllinux*"] \ No newline at end of file diff --git a/setup.py b/setup.py index b10a7e2e..507772d8 100644 --- a/setup.py +++ b/setup.py @@ -35,27 +35,9 @@ from setuptools import setup, find_packages from setuptools.command.install import install -# some code to handle the building of custom modules -FORCE_CUDA_EXTENSION = os.getenv("FORCE_CUDA_EXTENSION", "0") == "1" -BUILD_CPP = BUILD_CUDA = False - -# try to import torch -try: - import torch - - print(f"setup.py with torch {torch.__version__}") - from torch.utils.cpp_extension import BuildExtension, CppExtension - - BUILD_CPP = True - from torch.utils.cpp_extension import CUDA_HOME, CUDAExtension - - BUILD_CUDA = FORCE_CUDA_EXTENSION or (torch.cuda.is_available() and (CUDA_HOME is not None)) -except (ImportError, TypeError, AssertionError, AttributeError) as e: - warnings.warn(f"building custom extensions skipped: {e}") - def get_compile_args(module_name): """If user runs build with TORCH_HARMONICS_DEBUG=1 set, it will use debugging flags to build""" - + debug_mode = os.environ.get('TORCH_HARMONICS_DEBUG', '0') == '1' profile_mode = os.environ.get('TORCH_HARMONICS_PROFILE', '0') == '1' openmp_mode = os.getenv('TORCH_HARMONICS_ENABLE_OPENMP', '0') == '1' @@ -68,7 +50,7 @@ def get_compile_args(module_name): if profile_mode: nvcc_extra_flags.append("-lineinfo") nvcc_extra_flags.append("-Xptxas=-v") - + if debug_mode: print(f"WARNING: Compiling {module_name} with debugging flags") return { @@ -82,112 +64,161 @@ def get_compile_args(module_name): 'nvcc': ['-O3', "-DNDEBUG"] + nvcc_extra_flags } -def get_helpers_compile_args(): +def get_helpers_compile_args(BUILD_CPP, BUILD_CUDA): return { 'cxx': [ - f'-DBUILD_CPP={1 if BUILD_CPP else 0}', + f'-DBUILD_CPP={1 if BUILD_CPP else 0}', f'-DBUILD_CUDA={1 if BUILD_CUDA else 0}' - ], + ], } def get_ext_modules(): """Get list of extension modules to compile.""" - + + # some code to handle the building of custom modules + FORCE_CUDA_EXTENSION = os.getenv("FORCE_CUDA_EXTENSION", "0") == "1" + BUILD_CPP = BUILD_CUDA = False + + # PyTorch is required for building this package + try: + import torch + print(f"setup.py with torch {torch.__version__}") + from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDA_HOME, CUDAExtension + + print(f"Building with C++11 ABI = {torch._C._GLIBCXX_USE_CXX11_ABI}") + print(f"Compile flag will be -D_GLIBCXX_USE_CXX11_ABI={int(torch._C._GLIBCXX_USE_CXX11_ABI)}") + + BUILD_CPP = True + BUILD_CUDA = FORCE_CUDA_EXTENSION or (torch.cuda.is_available() and (CUDA_HOME is not None)) + + if BUILD_CUDA: + print("CUDA extensions will be built") + else: + print("CPU-only extensions will be built") + + except (ImportError, TypeError, AssertionError, AttributeError) as e: + raise RuntimeError(f"PyTorch is required to build torch-harmonics. Please install PyTorch first. Error: {e}") + ext_modules = [] cmdclass = {} + # Always build helper extensions (PyTorch is guaranteed to be available) print(f"Compiling helper routines for torch-harmonics.") ext_modules.append( CppExtension( - "disco_helpers", + "disco_helpers", [ "torch_harmonics/disco/csrc/disco_helpers.cpp", ], - extra_compile_args=get_helpers_compile_args(), + extra_compile_args=get_helpers_compile_args(BUILD_CPP, BUILD_CUDA), ) ) ext_modules.append( CppExtension( - "attention_helpers", + "attention_helpers", [ "torch_harmonics/attention/csrc/attention_helpers.cpp", ], - extra_compile_args=get_helpers_compile_args(), + extra_compile_args=get_helpers_compile_args(BUILD_CPP, BUILD_CUDA), ) ) - if BUILD_CPP: - # DISCO - # Create a single extension that includes both CPU and CUDA code - disco_sources = [ - "torch_harmonics/disco/csrc/disco_interface.cpp", - "torch_harmonics/disco/csrc/disco_cpu.cpp" - ] - - if BUILD_CUDA: - print(f"Compiling custom CUDA kernels for torch-harmonics.") - disco_sources.extend([ - "torch_harmonics/disco/csrc/disco_cuda_fwd.cu", - "torch_harmonics/disco/csrc/disco_cuda_bwd.cu", - ]) - ext_modules.append( - CUDAExtension( - "torch_harmonics.disco._C", - disco_sources, - extra_compile_args=get_compile_args("disco") - ) + # Always build main extensions + # DISCO + # Create a single extension that includes both CPU and CUDA code + disco_sources = [ + "torch_harmonics/disco/csrc/disco_interface.cpp", + "torch_harmonics/disco/csrc/disco_cpu.cpp" + ] + + if BUILD_CUDA: + print(f"Compiling custom CUDA kernels for torch-harmonics.") + disco_sources.extend([ + "torch_harmonics/disco/csrc/disco_cuda_fwd.cu", + "torch_harmonics/disco/csrc/disco_cuda_bwd.cu", + ]) + ext_modules.append( + CUDAExtension( + "torch_harmonics.disco._C", + disco_sources, + extra_compile_args=get_compile_args("disco") ) - else: - ext_modules.append( - CppExtension( - "torch_harmonics.disco._C", - disco_sources, - extra_compile_args=get_compile_args("disco") - ) + ) + else: + ext_modules.append( + CppExtension( + "torch_harmonics.disco._C", + disco_sources, + extra_compile_args=get_compile_args("disco") ) - cmdclass["build_ext"] = BuildExtension - - # ATTENTION - # Create a single extension that includes both CPU and CUDA code - attention_sources = [ - "torch_harmonics/attention/csrc/attention_interface.cpp", - "torch_harmonics/attention/csrc/attention_cpu_fwd.cpp", - "torch_harmonics/attention/csrc/attention_cpu_bwd.cpp", - ] - - if BUILD_CUDA: - print(f"Compiling attention CUDA kernels for torch-harmonics.") - attention_sources.extend([ - "torch_harmonics/attention/csrc/attention_cuda_utils.cu", - "torch_harmonics/attention/csrc/attention_cuda_fwd.cu", - "torch_harmonics/attention/csrc/attention_cuda_bwd.cu", - ]) - ext_modules.append( - CUDAExtension( - "torch_harmonics.attention._C", - attention_sources, - extra_compile_args=get_compile_args("attention") - ) + ) + cmdclass["build_ext"] = BuildExtension + + # ATTENTION + # Create a single extension that includes both CPU and CUDA code + attention_sources = [ + "torch_harmonics/attention/csrc/attention_interface.cpp", + "torch_harmonics/attention/csrc/attention_cpu_fwd.cpp", + "torch_harmonics/attention/csrc/attention_cpu_bwd.cpp", + ] + + if BUILD_CUDA: + print(f"Compiling attention CUDA kernels for torch-harmonics.") + attention_sources.extend([ + "torch_harmonics/attention/csrc/attention_cuda_utils.cu", + "torch_harmonics/attention/csrc/attention_cuda_fwd.cu", + "torch_harmonics/attention/csrc/attention_cuda_bwd.cu", + ]) + ext_modules.append( + CUDAExtension( + "torch_harmonics.attention._C", + attention_sources, + extra_compile_args=get_compile_args("attention") ) - else: - ext_modules.append( - CppExtension( - "torch_harmonics.attention._C", - attention_sources, - extra_compile_args=get_compile_args("attention") - ) + ) + else: + ext_modules.append( + CppExtension( + "torch_harmonics.attention._C", + attention_sources, + extra_compile_args=get_compile_args("attention") ) - cmdclass["build_ext"] = BuildExtension + ) + cmdclass["build_ext"] = BuildExtension return ext_modules, cmdclass if __name__ == "__main__": - ext_modules, cmdclass = get_ext_modules() + if any(cmd in sys.argv for cmd in ["build_ext", "bdist_wheel", "install", "develop"]): + ext_modules, cmdclass = get_ext_modules() + else: + ext_modules = [] + cmdclass = {} setup( + name="torch_harmonics", packages=find_packages(), ext_modules=ext_modules, cmdclass=cmdclass, + python_requires=">=3.9", + install_requires=[ + "torch>=2.4.0", + "numpy>=1.22.4", + ], + extras_require={ + "dev": [ + "pytest>=6.0.0", + "coverage>=6.5.0", + ], + "2d3ds": [ + "requests", + "tarfile", + "tqdm", + "PIL", + "h5py", + ], + }, + zip_safe=False, # Required for extensions ) diff --git a/torch_harmonics/__init__.py b/torch_harmonics/__init__.py index 2d390b54..d9a56242 100644 --- a/torch_harmonics/__init__.py +++ b/torch_harmonics/__init__.py @@ -29,7 +29,11 @@ # OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. # -__version__ = "0.8.1" +try: + from importlib.metadata import version, PackageNotFoundError + __version__ = version("torch_harmonics") +except (ImportError, PackageNotFoundError): + __version__ = "0.8.1" # fallback for development installs from .sht import RealSHT, InverseRealSHT, RealVectorSHT, InverseRealVectorSHT from .disco import DiscreteContinuousConvS2, DiscreteContinuousConvTransposeS2