diff --git a/.github/autoversion.py b/.github/autoversion.py deleted file mode 100644 index dbe884e6..00000000 --- a/.github/autoversion.py +++ /dev/null @@ -1,44 +0,0 @@ -""" -Sets the version in a pyproject.toml from git tags of the form -vX.Y.Z on HEAD. - -Usage: python3 .github/autoversion.py [...] -""" - -import re -import subprocess -import sys -from pathlib import Path - -import tomlkit - -TAG_PATTERN = re.compile(r"v([0-9]+\.[0-9]+\.[0-9]+)") - - -def get_version_from_tag(): - tags = subprocess.run( - ["git", "tag", "--points-at", "HEAD"], - check=True, - capture_output=True, - text=True, - ).stdout.splitlines() - - versions = [m.group(1) for m in map(TAG_PATTERN.fullmatch, tags) if m] - assert len(versions) == 1 - return versions[0] - - -def set_version(pyproject, version): - doc = tomlkit.parse(pyproject.read_text()) - assert "version" in doc["project"] - doc["project"]["version"] = version - pyproject.write_text(tomlkit.dumps(doc)) - - -if __name__ == "__main__": - if len(sys.argv) < 2: - sys.exit(__doc__) - - version = get_version_from_tag() - for path in sys.argv[1:]: - set_version(Path(path), version) diff --git a/.github/release.py b/.github/release.py new file mode 100644 index 00000000..d7a645b8 --- /dev/null +++ b/.github/release.py @@ -0,0 +1,164 @@ +""" +Release pipeline steps for .github/workflows/release.yaml. + + python3 .github/release.py resolve-symlinks + python3 .github/release.py set-version [...] + python3 .github/release.py install-cuda-toolkits [...] + python3 .github/release.py install-rocm-sdk +""" + +import argparse +import platform +import re +import shutil +import subprocess +import sys +import tempfile +import tomllib +from pathlib import Path + +TAG_PATTERN = re.compile(r"v([0-9]+\.[0-9]+\.[0-9]+)") +GIT_SYMLINK_MODE = "120000" +CUDA_ROOT_PATTERN = re.compile(r"/usr/local/cuda-([0-9]+)\.([0-9]+)") +CUDA_REPO_ARCH = {"x86_64": "x86_64", "aarch64": "sbsa"} +ROCM_INDEX_PATTERN = re.compile( + r"https://download\.pytorch\.org/whl/rocm[0-9]+\.[0-9]+" +) +ROCM_VERSION_PATTERN = re.compile(r"[0-9]+\.[0-9]+\.[0-9]+") +SHA256_PATTERN = re.compile(r"[0-9a-f]{64}") + + +def resolve_symlinks(args): + """Replaces every symlink tracked by git with a copy of its target.""" + entries = subprocess.run( + ["git", "ls-files", "-s", "-z"], + check=True, + capture_output=True, + text=True, + ).stdout.split("\0") + + for entry in filter(None, entries): + info, path = entry.split("\t", 1) + if info.split()[0] == GIT_SYMLINK_MODE: + link = Path(path) + target = link.resolve() + link.unlink() + if target.is_dir(): + shutil.copytree(target, link) + else: + shutil.copy2(target, link) + print(f"Resolved symlink {link} -> {target}") + + +def get_version_from_tag(): + tags = subprocess.run( + ["git", "tag", "--points-at", "HEAD"], + check=True, + capture_output=True, + text=True, + ).stdout.splitlines() + + versions = [m.group(1) for m in map(TAG_PATTERN.fullmatch, tags) if m] + assert len(versions) == 1 + return versions[0] + + +def set_version(args): + """Sets the version in each pyproject.toml from the git tag of the form + vX.Y.Z on HEAD.""" + import tomlkit + + version = get_version_from_tag() + for pyproject in args.pyprojects: + doc = tomlkit.parse(pyproject.read_text()) + assert "version" in doc["project"] + doc["project"]["version"] = version + pyproject.write_text(tomlkit.dumps(doc)) + + +def install_cuda_toolkits(args): + """Installs the CUDA headers and link-time libraries for each toolkit root + inside a manylinux container.""" + subprocess.run( + [ + "dnf", + "config-manager", + "--add-repo", + "https://developer.download.nvidia.com/compute/cuda/repos/rhel8/" + f"{CUDA_REPO_ARCH[platform.machine()]}/cuda-rhel8.repo", + ], + check=True, + ) + + for root in args.cuda_roots: + match = CUDA_ROOT_PATTERN.fullmatch(root) + assert match, f"expected a CUDA root like /usr/local/cuda-12.6, got {root!r}" + pkg_version = "-".join(match.groups()) + packages = [ + "cuda-nvcc", + "cuda-cudart-devel", + "cuda-nvrtc-devel", + "cuda-driver-devel", + ] + subprocess.run( + ["dnf", "install", "-y", *(f"{pkg}-{pkg_version}" for pkg in packages)], + check=True, + ) + + +def install_rocm_sdk(args): + """Installs rocm-sdk-core (HIP headers, libamdhip64 and libhiprtc), pinned + in [tool.openequivariance.rocm-sdk-core] of pyproject.toml, into a target + directory; skipped on non-x86_64 machines.""" + if platform.machine() == "x86_64": + with open(args.pyproject, "rb") as f: + pin = tomllib.load(f)["tool"]["openequivariance"]["rocm-sdk-core"] + assert ROCM_INDEX_PATTERN.fullmatch(pin["index-url"]) + assert ROCM_VERSION_PATTERN.fullmatch(pin["version"]) + assert SHA256_PATTERN.fullmatch(pin["sha256"]) + with tempfile.TemporaryDirectory() as tmp: + requirements = Path(tmp) / "rocm_requirements.txt" + requirements.write_text( + f"--index-url {pin['index-url']}\n" + f"rocm-sdk-core=={pin['version']} " + f"--hash=sha256:{pin['sha256']}\n" + ) + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "install", + "--no-deps", + "--require-hashes", + "--target", + str(args.target), + "-r", + str(requirements), + ], + check=True, + ) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0].strip()) + steps = parser.add_subparsers(dest="step", required=True) + + step = steps.add_parser("resolve-symlinks", help=resolve_symlinks.__doc__) + step.set_defaults(func=resolve_symlinks) + + step = steps.add_parser("set-version", help=set_version.__doc__) + step.add_argument("pyprojects", nargs="+", type=Path) + step.set_defaults(func=set_version) + + step = steps.add_parser("install-cuda-toolkits", help=install_cuda_toolkits.__doc__) + step.add_argument("cuda_roots", nargs="+") + step.set_defaults(func=install_cuda_toolkits) + + step = steps.add_parser("install-rocm-sdk", help=install_rocm_sdk.__doc__) + step.add_argument("pyproject", type=Path) + step.add_argument("target", type=Path) + step.set_defaults(func=install_rocm_sdk) + + args = parser.parse_args() + args.func(args) diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index b5fe3451..7f5afd7d 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -4,8 +4,41 @@ on: workflow_dispatch: jobs: - build-oeq: - name: Build oeq dist + build-oeq-wheels: + name: Build oeq wheels (${{ matrix.arch }}) + runs-on: ${{ matrix.runner }} + timeout-minutes: 60 + strategy: + matrix: + include: + - runner: ubuntu-24.04 + arch: x86_64 + - runner: ubuntu-24.04-arm + arch: aarch64 + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + fetch-tags: true + - name: set up Python + uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 + with: + python-version: '3.12' + - name: replace symlinks with copies + run: python3 .github/release.py resolve-symlinks + - name: set version from tag + run: | + python3 -m pip install -r .github/workflows/requirements_release_ci.txt + python3 .github/release.py set-version openequivariance/pyproject.toml + - name: build wheels + run: python3 -m cibuildwheel --output-dir wheelhouse openequivariance + - name: store the wheels + uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 + with: + name: oeq-wheels-${{ matrix.arch }} + path: wheelhouse/*.whl + + build-oeq-sdist: + name: Build oeq sdist runs-on: ubuntu-latest steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 @@ -14,27 +47,29 @@ jobs: - name: set up Python uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: - python-version: '3.10' + python-version: '3.12' + - name: replace symlinks with copies + run: python3 .github/release.py resolve-symlinks - name: set version from tag run: | python3 -m pip install -r .github/workflows/requirements_release_ci.txt - python3 .github/autoversion.py openequivariance/pyproject.toml + python3 .github/release.py set-version openequivariance/pyproject.toml - name: install dependencies, then build source tarball run: | cd openequivariance - python3 -m pip install build + python3 -m pip install build python3 -m build --sdist - - name: store the distribution packages + - name: store the distribution packages uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 with: - name: oeq-package-distributions + name: oeq-sdist path: openequivariance/dist/ pypi-publish: name: Upload oeq to PyPI runs-on: ubuntu-latest - needs: build-oeq - environment: + needs: [build-oeq-wheels, build-oeq-sdist] + environment: name: pypi url: https://pypi.org/p/openequivariance permissions: @@ -43,8 +78,9 @@ jobs: - name: download the distributions uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2 with: - name: oeq-package-distributions - path: dist/ + pattern: oeq-* + merge-multiple: true + path: dist/ - name: publish package distributions to PyPI uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3 @@ -61,16 +97,18 @@ jobs: uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: python-version: '3.10' + - name: replace symlinks with copies + run: python3 .github/release.py resolve-symlinks - name: set version from tag run: | python3 -m pip install -r .github/workflows/requirements_release_ci.txt - python3 .github/autoversion.py openequivariance_extjax/pyproject.toml + python3 .github/release.py set-version openequivariance_extjax/pyproject.toml - name: install dependencies, then build source tarball run: | cd openequivariance_extjax - python3 -m pip install build - python3 -m build --sdist - - name: store the distribution packages + python3 -m pip install build + python3 -m build --sdist + - name: store the distribution packages uses: actions/upload-artifact@cf430e030ddbb5b0abf93d22962f4752f3646cd9 # v7.0.2 with: name: oeqjax-package-distributions @@ -80,7 +118,7 @@ jobs: name: Upload extjax to PyPI runs-on: ubuntu-latest needs: build-oeq-extjax - environment: + environment: name: pypi url: https://pypi.org/p/openequivariance_extjax permissions: @@ -90,6 +128,6 @@ jobs: uses: actions/download-artifact@9000827ccba6bdab643e8b6fd33ac0654aef8333 # v8.0.2 with: name: oeqjax-package-distributions - path: dist/ + path: dist/ - name: publish package distributions to PyPI uses: pypa/gh-action-pypi-publish@7245906fb61fa33a13554f9d5fb856fe61758086 # v1.14.3 diff --git a/.github/workflows/requirements_cuda_ci.txt b/.github/workflows/requirements_cuda_ci.txt index 4872ba8d..09f2376a 100644 --- a/.github/workflows/requirements_cuda_ci.txt +++ b/.github/workflows/requirements_cuda_ci.txt @@ -1,6 +1,7 @@ +--extra-index-url https://download.pytorch.org/whl/cu129 numpy==2.2.5 -torch==2.13.0 --index-url https://download.pytorch.org/whl/cu129 +torch==2.13.0+cu129 pytest==9.0.3 ninja==1.11.1.4 -nanobind==2.10.2 +nanobind==3.1.0 scikit-build-core==0.11.6 \ No newline at end of file diff --git a/.github/workflows/requirements_release_ci.txt b/.github/workflows/requirements_release_ci.txt index 7b69eff7..dcb25b15 100644 --- a/.github/workflows/requirements_release_ci.txt +++ b/.github/workflows/requirements_release_ci.txt @@ -1 +1,2 @@ tomlkit==0.15.2 +cibuildwheel==4.3.5 diff --git a/.github/workflows/requirements_sycl_ci.txt b/.github/workflows/requirements_sycl_ci.txt index d2d92ac3..a4b9ca6e 100644 --- a/.github/workflows/requirements_sycl_ci.txt +++ b/.github/workflows/requirements_sycl_ci.txt @@ -3,5 +3,5 @@ numpy==2.2.5 torch==2.13.0+xpu pytest==9.0.3 ninja==1.11.1.4 -nanobind==2.10.2 +nanobind==3.1.0 scikit-build-core==0.11.6 diff --git a/.github/workflows/verify_extension_build.yml b/.github/workflows/verify_extension_build.yml index c6b7c16a..c33a4da2 100644 --- a/.github/workflows/verify_extension_build.yml +++ b/.github/workflows/verify_extension_build.yml @@ -20,7 +20,7 @@ jobs: - name: Set up Python uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: - python-version: "3.10" + python-version: "3.12" cache: 'pip' cache-dependency-path: '**/requirements_cuda_ci.txt' - name: Install dependencies @@ -35,14 +35,31 @@ jobs: run: | pytest tests/import_test.py - export OEQ_JIT_EXTENSION=1 - - pytest tests/import_test.py - - name: Test JAX extension build run: | XLA_DIRECT_DOWNLOAD=1 pip install -e "./openequivariance_extjax" --no-build-isolation + verify_hip_extension: + if: ${{ github.event.label.name == 'ci-ready' || github.event_name != 'pull_request' }} + runs-on: ubuntu-latest + + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Set up Python + uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 + with: + python-version: "3.12" + + - name: Install pinned rocm-sdk-core + run: | + python .github/release.py install-rocm-sdk openequivariance/pyproject.toml "$RUNNER_TEMP/rocm-sdk-core" + + - name: Build HIP extension + env: + ROCM_PATH: ${{ runner.temp }}/rocm-sdk-core/_rocm_sdk_core + run: | + pip wheel --no-deps -w dist ./openequivariance + verify_sycl_extension: if: ${{ github.event.label.name == 'ci-ready' || github.event_name != 'pull_request' }} runs-on: ubuntu-latest @@ -52,7 +69,7 @@ jobs: - name: Set up Python uses: actions/setup-python@6d4fd22b20465ccc0f7354476588e8df7c9ca0e0 # v7.1.0 with: - python-version: "3.10" + python-version: "3.12" cache: 'pip' cache-dependency-path: '**/requirements_sycl_ci.txt' @@ -70,7 +87,9 @@ jobs: run: | python -m pip install --upgrade pip pip install -r .github/workflows/requirements_sycl_ci.txt - pip install -e "./openequivariance" + source /opt/intel/oneapi/setvars.sh + export CXX=icpx + pip install -e "./openequivariance[sycl]" - name: Test SYCL extension build via import run: | @@ -78,7 +97,3 @@ jobs: export CXX=icpx pytest tests/import_test.py - - export OEQ_JIT_EXTENSION=1 - - pytest tests/import_test.py \ No newline at end of file diff --git a/README.md b/README.md index 34b88cee..24fa7a0f 100644 --- a/README.md +++ b/README.md @@ -13,8 +13,7 @@ that [e3nn](https://e3nn.org/) supports commonly found in graph neural networks (e.g. [Nequip](https://github.com/mir-group/nequip) or [MACE](https://github.com/ACEsuit/mace)). To get -started with PyTorch, ensure that you have PyTorch -and GCC 9+ available before installing our package via +started, make sure that you have PyTorch installed and run: ```bash pip install openequivariance diff --git a/docs/installation.rst b/docs/installation.rst index 17feab14..8405f77f 100644 --- a/docs/installation.rst +++ b/docs/installation.rst @@ -8,11 +8,11 @@ Installation You need the following to install OpenEquivariance: - A Linux system equipped with an NVIDIA / AMD / Intel graphics card. -- Either PyTorch >= 2.4 (>= 2.10 for AOTI, export, and Intel GPUs), or - JAX>0.5.0 with CUDA, RocM, or XPU support. -- GCC 9+ and the CUDA / HIP toolkit, or the oneAPI DPC++ compiler - (``icpx``) for Intel GPUs. The command - ``c++ --version`` should return >= 9.0; see below for details on +- Either PyTorch >= 2.10, or + JAX>0.5.0 with CUDA, RocM, or XPU support. +- To build from source: GCC 9+ and the CUDA / HIP toolkit, or the + oneAPI DPC++ compiler (``icpx``) for Intel GPUs. The command + ``c++ --version`` should return >= 9.0; see below for details on setting an alternate compiler. .. note:: @@ -32,32 +32,30 @@ You need the following to install OpenEquivariance: pip install openequivariance python -c "import openequivariance" - The second line triggers a build of the C++ extension we use to compile - kernels, which can take a couple of minutes. Subsequent imports are - much faster since this extension is cached. + OpenEquivariance bundles a precompiled library files for two recent + CUDA versions and a HIP version. If you're building for an architecture + we don't support or need a different CUDA version, you'll need to build + the package from source: - To support ``torch.compile``, ``torch.export``, and - JITScript, OpenEquivariance needs to compile a C++ extension - tightly integrated with PyTorch. If you see a warning that - this extension could not be compiled, first check: + .. code-block:: bash - .. code-block:: bash + pip install --no-binary openequivariance openequivariance - c++ --version - - To build the extension with an alternate compiler, set the - ``CC`` and ``CXX`` - environment variable and retry the import: + Set the environment variables ``CC`` and ``CXX`` before the pip install + to build with a different compiler: .. code-block:: bash export CC=/path/to/your/gcc export CXX=/path/to/your/g++ - python -c "import openequivariance" + On Intel GPUs, the extension is still compiled on first import, + which can take a couple of minutes; subsequent imports use the + cached build. Install the extra build dependencies with: - These configuration steps are required only ONCE after - installation (or upgrade) with pip. + .. code-block:: bash + + pip install openequivariance[sycl] .. tab:: JAX NVIDIA GPUs @@ -93,11 +91,11 @@ You need the following to install OpenEquivariance: .. code-block:: bash - pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance[jax]" - pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax --no-build-isolation" + pip install "openequivariance[jax] @ git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance" + pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax" --no-build-isolation # Use the command below for JAX+AMD - # JAX_HIP=1 pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax --no-build-isolation" + # JAX_HIP=1 pip install "git+https://github.com/PASSIONLab/OpenEquivariance#subdirectory=openequivariance_extjax" --no-build-isolation If you're using JAX, set the environment variable @@ -116,45 +114,6 @@ Here are some tested environment configuration files. If you use OpenEquivarianc on a major cluster, send us a pull request to add your configuration! -.. tab:: NERSC Perlmutter (NVIDIA A100) - - .. code-block:: bash - :caption: env.sh (last updated June 2025) - - module load gcc - module load conda - - # Deactivate any base environments - for i in $(seq ${CONDA_SHLVL}); do - conda deactivate - done - - conda activate - - -.. tab:: OLCF Frontier (AMD MI250x) - - You need to install a HIP-enabled verison of PyTorch to use our package. - Follow the steps `here `_. - - - .. code-block:: bash - :caption: env.sh (last updated June 2025) - - module load PrgEnv-gnu/8.6.0 - module load miniforge3/23.11.0-0 - module load rocm/6.4.0 - module load craype-accel-amd-gfx90a - - for i in $(seq ${CONDA_SHLVL}); do - conda deactivate - done - - conda activate - export CC=cc - export CXX=CC - - .. tab:: ALCF Aurora (Intel Data Center GPU Max) Aurora provides a PyTorch build with XPU support in its ``frameworks`` diff --git a/openequivariance/CMakeLists.txt b/openequivariance/CMakeLists.txt index d1445395..1ceeadbd 100644 --- a/openequivariance/CMakeLists.txt +++ b/openequivariance/CMakeLists.txt @@ -1,22 +1,29 @@ cmake_minimum_required(VERSION 3.15...3.30) project(openequivariance_stable_ext) -find_package(Python 3.10 REQUIRED COMPONENTS Interpreter Development.Module) +find_package(Python 3.11 REQUIRED COMPONENTS Interpreter Development.Module ${SKBUILD_SABI_COMPONENT}) # Download LibTorch include(FetchContent) if (AARCH64 OR CMAKE_SYSTEM_PROCESSOR MATCHES "(arm)|(ARM)|(aarch64)") message(STATUS "Targeting ARM architecture (AARCH64 override or processor match)") - set(LIBTORCH_URL "https://github.com/PASSIONLab/libtorch_aarch64/raw/refs/heads/main/pkg-libtorch.tar.zst") + set(LIBTORCH_ARCH AARCH64) else() message(STATUS "Did not match arm/aarch64, performing x86 install") - set(LIBTORCH_URL "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.10.0%2Bcpu.zip") + set(LIBTORCH_ARCH X86_64) +endif() + +set(LIBTORCH_URL "${LIBTORCH_URL_${LIBTORCH_ARCH}}") +set(LIBTORCH_SHA256 "${LIBTORCH_SHA256_${LIBTORCH_ARCH}}") +if(NOT LIBTORCH_URL OR NOT LIBTORCH_SHA256) + message(FATAL_ERROR "LIBTORCH_URL / LIBTORCH_SHA256 not set") endif() FetchContent_Declare( libtorch URL "${LIBTORCH_URL}" + URL_HASH SHA256=${LIBTORCH_SHA256} ) message(STATUS "Downloading LibTorch...") @@ -52,11 +59,10 @@ set(OEQ_SOURCES ${EXT_JSON_DIR}/json11.cpp ) -set(OEQ_INSTALL_DIR "${CMAKE_CURRENT_SOURCE_DIR}/openequivariance/_torch/extlib") +set(OEQ_INSTALL_DIR "openequivariance/_torch/extlib") function(add_stable_extension target_name backend_define link_libraries) - # Create nanobind extension - nanobind_add_module(${target_name} NB_STATIC ${OEQ_SOURCES}) + nanobind_add_module(${target_name} BACKEND_MODULE nanobind_backend ${OEQ_SOURCES}) set_target_properties(${target_name} PROPERTIES CXX_STANDARD 17 @@ -65,10 +71,11 @@ function(add_stable_extension target_name backend_define link_libraries) ) # Enforce CXX11 ABI to match LibTorch - target_compile_definitions(${target_name} PRIVATE + target_compile_definitions(${target_name} PRIVATE ${backend_define}=1 _GLIBCXX_USE_CXX11_ABI=1 INCLUDE_NB_EXTENSION + EXTENSION_NAME=${target_name} ) target_include_directories(${target_name} PRIVATE @@ -115,34 +122,86 @@ function(add_stable_extension target_name backend_define link_libraries) install(TARGETS ${aoti_target_name} LIBRARY DESTINATION "${OEQ_INSTALL_DIR}") endfunction() -find_package(CUDAToolkit QUIET) find_package(hip QUIET) find_package(hiprtc QUIET) -if(CUDAToolkit_FOUND) - message(STATUS "Building stable extension with CUDA backend.") +if(NOT hip_FOUND AND DEFINED ENV{ROCM_PATH}) + find_path(HIP_INCLUDE_DIR hip/hip_runtime.h PATHS $ENV{ROCM_PATH}/include NO_DEFAULT_PATH) + find_library(AMDHIP64_LIB NAMES amdhip64 libamdhip64.so.7 PATHS $ENV{ROCM_PATH}/lib NO_DEFAULT_PATH) + find_library(HIPRTC_LIB NAMES hiprtc libhiprtc.so.7 PATHS $ENV{ROCM_PATH}/lib NO_DEFAULT_PATH) + + if(HIP_INCLUDE_DIR AND AMDHIP64_LIB AND HIPRTC_LIB) + add_library(hip::host INTERFACE IMPORTED) + target_include_directories(hip::host INTERFACE ${HIP_INCLUDE_DIR}) + target_compile_definitions(hip::host INTERFACE __HIP_PLATFORM_AMD__=1) + target_link_libraries(hip::host INTERFACE ${AMDHIP64_LIB}) + + add_library(hiprtc::hiprtc INTERFACE IMPORTED) + target_link_libraries(hiprtc::hiprtc INTERFACE ${HIPRTC_LIB}) + + set(hip_FOUND TRUE) + endif() +endif() + +set(OEQ_CUDA_ROOTS "$ENV{OEQ_CUDA_ROOTS}" CACHE STRING "CUDA toolkit roots to build against") +string(REPLACE " " ";" OEQ_CUDA_ROOTS "${OEQ_CUDA_ROOTS}") +set(OEQ_CUDA_FOUND FALSE) +if(OEQ_CUDA_ROOTS) + set(OEQ_CUDA_FOUND TRUE) +else() + find_package(CUDAToolkit QUIET) + set(OEQ_CUDA_FOUND ${CUDAToolkit_FOUND}) +endif() + +if(OEQ_CUDA_FOUND) add_library(cuda_stub_lib SHARED ${EXT_DIR}/stubs/stream.cpp) target_include_directories(cuda_stub_lib PRIVATE ${LIBTORCH_INCLUDE_DIR} ) - set_target_properties(cuda_stub_lib PROPERTIES - OUTPUT_NAME "torch_cuda" + set_target_properties(cuda_stub_lib PROPERTIES + OUTPUT_NAME "torch_cuda" POSITION_INDEPENDENT_CODE ON CXX_STANDARD 17 ) +endif() +if(OEQ_CUDA_FOUND AND NOT OEQ_CUDA_ROOTS) + message(STATUS "Building stable extension with CUDA ${CUDAToolkit_VERSION} backend.") set(CUDA_LINK_LIBS CUDA::cudart CUDA::cuda_driver - CUDA::nvrtc + ${CUDA_nvrtc_LIBRARY} cuda_stub_lib ) - add_stable_extension(oeq_stable_cuda CUDA_BACKEND "${CUDA_LINK_LIBS}") + add_stable_extension(oeq_stable_cuda${CUDAToolkit_VERSION_MAJOR} CUDA_BACKEND "${CUDA_LINK_LIBS}") endif() +foreach(cuda_root ${OEQ_CUDA_ROOTS}) + file(STRINGS "${cuda_root}/include/cuda.h" cuda_version_line REGEX "^#define CUDA_VERSION [0-9]+") + if(NOT cuda_version_line) + message(FATAL_ERROR "No CUDA toolkit found at ${cuda_root}") + endif() + string(REGEX REPLACE ".* ([0-9]+)$" "\\1" cuda_version "${cuda_version_line}") + math(EXPR cuda_major "${cuda_version} / 1000") + math(EXPR cuda_minor "(${cuda_version} % 1000) / 10") + message(STATUS "Building stable extension with CUDA ${cuda_major}.${cuda_minor} backend from ${cuda_root}.") + + set(cuda_lib_dir "${cuda_root}/lib64") + add_library(oeq_cuda${cuda_major}_toolkit INTERFACE) + target_include_directories(oeq_cuda${cuda_major}_toolkit INTERFACE "${cuda_root}/include") + target_link_libraries(oeq_cuda${cuda_major}_toolkit INTERFACE + "${cuda_lib_dir}/libcudart.so" + "${cuda_lib_dir}/libnvrtc.so" + "${cuda_lib_dir}/stubs/libcuda.so" + ) + + set(CUDA_LINK_LIBS oeq_cuda${cuda_major}_toolkit cuda_stub_lib) + add_stable_extension(oeq_stable_cuda${cuda_major} CUDA_BACKEND "${CUDA_LINK_LIBS}") +endforeach() + if(hip_FOUND) message(STATUS "Building stable extension with HIP backend.") @@ -200,6 +259,6 @@ if(OEQ_SYCL_FOUND) endforeach() endif() -if(NOT CUDAToolkit_FOUND AND NOT hip_FOUND AND NOT OEQ_SYCL_FOUND) - message(WARNING "None of CUDAToolkit, HIP or SYCL was found. The stable extension will not be built.") +if(NOT OEQ_CUDA_FOUND AND NOT hip_FOUND AND NOT OEQ_SYCL_FOUND) + message(FATAL_ERROR "None of CUDAToolkit, HIP or SYCL was found") endif() diff --git a/openequivariance/README.md b/openequivariance/README.md deleted file mode 100644 index 3efe59a7..00000000 --- a/openequivariance/README.md +++ /dev/null @@ -1,220 +0,0 @@ -# OpenEquivariance -[![OEQ C++ Extension Build Verification](https://github.com/PASSIONLab/OpenEquivariance/actions/workflows/verify_extension_build.yml/badge.svg?event=push)](https://github.com/PASSIONLab/OpenEquivariance/actions/workflows/verify_extension_build.yml) -[![License](https://img.shields.io/badge/License-BSD_3--Clause-blue.svg)](https://opensource.org/licenses/BSD-3-Clause) - -[[PyTorch Examples]](#pytorch-examples) -[[JAX Examples]](#jax-examples) -[[Citation and Acknowledgements]](#citation-and-acknowledgements) - -OpenEquivariance is a CUDA and HIP kernel generator for the Clebsch-Gordon tensor product, -a key kernel in rotation-equivariant deep neural networks. -It implements some of the tensor products -that [e3nn](https://e3nn.org/) supports -commonly found in graph neural networks -(e.g. [Nequip](https://github.com/mir-group/nequip) or -[MACE](https://github.com/ACEsuit/mace)). To get -started with PyTorch, ensure that you have PyTorch -and GCC 9+ available before installing our package via - -```bash -pip install openequivariance -``` - -We provide up to an order of magnitude acceleration over e3nn perform on par with the latest -version of [NVIDIA cuEquivariance](https://github.com/NVIDIA/cuEquivariance), -which has a closed-source kernel package. -We also offer fused equivariant graph -convolutions that can reduce -computation and memory consumption significantly. - -For detailed instructions on tests, benchmarks, MACE / Nequip, and our API, -check out the [documentation](https://passionlab.github.io/OpenEquivariance). - -⭐️ **JAX**: Our latest update brings -support for JAX. For NVIDIA GPUs, -install it (after installing JAX) -with the following two commands strictly in order: - -``` bash -pip install openequivariance[jax] -pip install openequivariance_extjax --no-build-isolation -``` - -For AMD GPUs: -``` bash -pip install openequivariance[jax] -JAX_HIP=1 pip install openequivariance_extjax --no-build-isolation -``` - -See the section below for example usage and -our [API page](https://passionlab.github.io/OpenEquivariance/api/) for more details. - -## PyTorch Examples -Here's a CG tensor product implemented by e3nn: - -```python -import torch -import e3nn.o3 as o3 - -gen = torch.Generator(device='cuda') - -batch_size = 1000 -X_ir, Y_ir, Z_ir = o3.Irreps("1x2e"), o3.Irreps("1x3e"), o3.Irreps("1x2e") -X = torch.rand(batch_size, X_ir.dim, device='cuda', generator=gen) -Y = torch.rand(batch_size, Y_ir.dim, device='cuda', generator=gen) - -instructions=[(0, 0, 0, "uvu", True)] - -tp_e3nn = o3.TensorProduct(X_ir, Y_ir, Z_ir, instructions, - shared_weights=False, internal_weights=False).to('cuda') -W = torch.rand(batch_size, tp_e3nn.weight_numel, device='cuda', generator=gen) - -Z = tp_e3nn(X, Y, W) -print(torch.norm(Z)) -``` - -And here's the same tensor product using openequivariance. We require that your -tensors are stored on a CUDA device for this to work: - -```python -import openequivariance as oeq - -problem = oeq.TPProblem(X_ir, Y_ir, Z_ir, instructions, shared_weights=False, internal_weights=False) -tp_fast = oeq.TensorProduct(problem, torch_op=True) - -Z = tp_fast(X, Y, W) # Reuse X, Y, W from earlier -print(torch.norm(Z)) -``` - -Our interface for `oeq.TPProblem` is almost a strict superset of -`o3.TensorProduct` (two key differences: we -impose `internal_weights=False` and add support for multiple datatypes). -You can pass e3nn `Irreps` instances directly or -use `oeq.Irreps`, which is identical. - -We recommend reading the [e3nn documentation and API reference](https://docs.e3nn.org/en/latest/) first, then using our kernels -as drop-in replacements. We support most "uvu" and "uvw" tensor products; -see [this section](#tensor-products-we-accelerate) for an up-to-date list of supported configurations. - -**Important**: For many configurations, our kernels return results identical to -e3nn up to floating point roundoff (this includes all "uvu" problems with -multiplicity 1 for all irreps in the second input). For other configurations -(e.g. any "uvw" connection modes), we return identical -results up to a well-defined reordering of the weights relative to e3nn. - -If you're executing tensor products as part of a message passing graph -neural network, we offer fused kernels that save both memory and compute time: - -```python -from torch_geometric import EdgeIndex - -node_ct, nonzero_ct = 3, 4 - -# Receiver, sender indices for message passing GNN -edge_index = EdgeIndex( - [[0, 1, 1, 2], # Receiver - [1, 0, 2, 1]], # Sender - device='cuda', - dtype=torch.long) - -X = torch.rand(node_ct, X_ir.dim, device='cuda', generator=gen) -Y = torch.rand(nonzero_ct, Y_ir.dim, device='cuda', generator=gen) -W = torch.rand(nonzero_ct, problem.weight_numel, device='cuda', generator=gen) - -tp_conv = oeq.TensorProductConv(problem, torch_op=True, deterministic=False) # Reuse problem from earlier -Z = tp_conv.forward(X, Y, W, edge_index[0], edge_index[1]) # Z has shape [node_ct, z_ir.dim] -print(torch.norm(Z)) -``` - -If you can guarantee `EdgeIndex` is sorted by receiver index and supply the transpose -permutation, we can provide even greater speedup (and deterministic results) -by avoiding atomics: - -```python -_, sender_perm = edge_index.sort_by("col") # Sort by sender index -edge_index, receiver_perm = edge_index.sort_by("row") # Sort by receiver index - -# Now we can use the faster deterministic algorithm -tp_conv = oeq.TensorProductConv(problem, torch_op=True, deterministic=True) -Z = tp_conv.forward(X, Y[receiver_perm], W[receiver_perm], edge_index[0], edge_index[1], sender_perm) -print(torch.norm(Z)) -``` -**Note**: you don't need Pytorch geometric to use our kernels. When -`deterministic=False`, the `sender` and `receiver` indices can have -arbitrary order. - -## JAX Examples -After installation, use the library -as follows. Set `OEQ_NOTORCH=1` -in your environment to avoid the PyTorch import in -the regular `openequivariance` package. -```python -import jax -import os - -os.environ["OEQ_NOTORCH"] = "1" -import openequivariance as oeq - -seed = 42 -key = jax.random.PRNGKey(seed) - -batch_size = 1000 -X_ir, Y_ir, Z_ir = oeq.Irreps("1x2e"), oeq.Irreps("1x3e"), oeq.Irreps("1x2e") -problem = oeq.TPProblem(X_ir, Y_ir, Z_ir, [(0, 0, 0, "uvu", True)], shared_weights=False, internal_weights=False) - - -node_ct, nonzero_ct = 3, 4 -edge_index = jax.numpy.array( - [ - [0, 1, 1, 2], - [1, 0, 2, 1], - ], - dtype=jax.numpy.int32, # NOTE: This int32, not int64 -) - -X = jax.random.uniform(key, shape=(node_ct, X_ir.dim), minval=0.0, maxval=1.0, dtype=jax.numpy.float32) -Y = jax.random.uniform(key, shape=(nonzero_ct, Y_ir.dim), - minval=0.0, maxval=1.0, dtype=jax.numpy.float32) -W = jax.random.uniform(key, shape=(nonzero_ct, problem.weight_numel), - minval=0.0, maxval=1.0, dtype=jax.numpy.float32) - -tp_conv = oeq.jax.TensorProductConv(problem, deterministic=False) -Z = tp_conv.forward( - X, Y, W, edge_index[0], edge_index[1] -) -print(jax.numpy.linalg.norm(Z)) - -# Test JAX JIT -jitted = jax.jit(lambda X, Y, W, e1, e2: tp_conv.forward(X, Y, W, e1, e2)) -print(jax.numpy.linalg.norm(jitted(X, Y, W, edge_index[0], edge_index[1]))) -``` - -## Citation and Acknowledgements -If you find this code useful, please cite our paper: - -```bibtex -@inbook{openequivariance, -author={Vivek Bharadwaj and Austin Glover and Aydin Buluc and James Demmel}, -title={An Efficient Sparse Kernel Generator for O(3)-Equivariant Deep Networks}, -booktitle = {SIAM Conference on Applied and Computational Discrete Algorithms (ACDA25)}, -chapter = {}, -url={https://arxiv.org/abs/2501.13986}, -publisher={Society for Industrial and Applied Mathematics}, -year={2025} -} -``` - -Our codebase includes a lightweight clone of -[e3nn](https://e3nn.org/)'s frontend interface (in particular, the -`TensorProduct` and `Irreps` classes). We removed references to Pytorch -and separated the implementation from the problem description (for future -frontend support outside of torch). We also extracted the Wigner 3j tensor generating code from QuTiP. Thank you to the current -developers and maintainers! - -## Copyright - -Copyright (c) 2025, The Regents of the University of California, through Lawrence Berkeley National Laboratory (subject to receipt of any required approvals from the U.S. Dept. of Energy). All rights reserved. - -If you have questions about your rights to use or distribute this software, please contact Berkeley Lab's Intellectual Property Office at IPO@lbl.gov. - -NOTICE. This Software was developed under funding from the U.S. Department of Energy and the U.S. Government consequently retains certain rights. As such, the U.S. Government has been granted for itself and others acting on its behalf a paid-up, nonexclusive, irrevocable, worldwide license in the Software to reproduce, distribute copies to the public, prepare derivative works, and perform publicly and display publicly, and to permit others to do so. \ No newline at end of file diff --git a/openequivariance/README.md b/openequivariance/README.md new file mode 120000 index 00000000..32d46ee8 --- /dev/null +++ b/openequivariance/README.md @@ -0,0 +1 @@ +../README.md \ No newline at end of file diff --git a/openequivariance/openequivariance/__init__.py b/openequivariance/openequivariance/__init__.py index 1a3681f5..b8e02d74 100644 --- a/openequivariance/openequivariance/__init__.py +++ b/openequivariance/openequivariance/__init__.py @@ -58,19 +58,15 @@ def _check_package_editable(): ] ) - from openequivariance._torch.extlib import ( - LINKED_LIBPYTHON, - LINKED_LIBPYTHON_ERROR, - BUILT_EXTENSION, - BUILT_EXTENSION_ERROR, - USE_PRECOMPILED_EXTENSION, - ) + from openequivariance._torch.extlib import BUILT_EXTENSION, BUILT_EXTENSION_ERROR def torch_ext_so_path(): """ :returns: Path to a ``.so`` file that must be linked to use OpenEquivariance - from the PyTorch C++ Interface. + from the PyTorch C++ Interface, or ``None`` if PyTorch support is + disabled (``OEQ_NOTORCH=1``). + :raises ImportError: If the OpenEquivariance extension failed to load. """ try: return torch_ext_so_path_internal() @@ -83,14 +79,6 @@ def torch_ext_so_path(): import openequivariance_extjax import openequivariance.jax as jax - # TODO-someday: enable - # extjax_version = version("openequivariance_extjax") - # if extjax_version != __version__: - # warnings.warn( - # f"openequivariance_extjax version {extjax_version} does not match " - # f"openequivariance version {__version__}. Ensure both versions match." - # ) - except Exception as e: error = e diff --git a/openequivariance/openequivariance/_torch/extlib/__init__.py b/openequivariance/openequivariance/_torch/extlib/__init__.py index 4ec88480..54f31155 100644 --- a/openequivariance/openequivariance/_torch/extlib/__init__.py +++ b/openequivariance/openequivariance/_torch/extlib/__init__.py @@ -5,161 +5,112 @@ import warnings import sysconfig import contextlib +import importlib +import importlib.util from pathlib import Path from packaging.version import Version import torch -from openequivariance.core.logging import getLogger - -oeq_root = str(Path(__file__).parent.parent.parent) - -BUILT_EXTENSION = False -BUILT_EXTENSION_ERROR = None - -LINKED_LIBPYTHON = False -LINKED_LIBPYTHON_ERROR = None - -extension_module = None - - -def _detect_backend(): - """HIP builds report a ``torch.version.cuda`` of None, so HIP is tested first.""" - if torch.version.hip: - return "hip" - if torch.version.cuda: - return "cuda" - if getattr(torch.version, "xpu", None): - return "sycl" - return None - - -BACKEND = _detect_backend() - -assert BACKEND is not None, ( - "Only the CUDA, HIP and XPU (SYCL) backends are supported. " - "No supported accelerator was detected in this PyTorch build." -) - -IS_HIP = BACKEND == "hip" -IS_SYCL = BACKEND == "sycl" - -if IS_SYCL and Version(torch.__version__) < Version("2.10"): +if Version(torch.__version__) < Version("2.10"): raise RuntimeError( - f"The SYCL backend requires PyTorch >= 2.10, found {torch.__version__}." + f"OpenEquivariance requires PyTorch >= 2.10, found {torch.__version__}." ) +oeq_root = str(Path(__file__).parent.parent.parent) -DEVICE_TYPE = "xpu" if IS_SYCL else "cuda" - - -@contextlib.contextmanager -def _maybe_set_max_jobs_env(limit=8): - if "MAX_JOBS" in os.environ: - yield - return - os.environ["MAX_JOBS"] = str(min(limit, os.cpu_count() or 1)) - try: - yield - finally: - os.environ.pop("MAX_JOBS", None) - - -def _wait_for_torch_build_lock(name, timeout=300): - from torch.utils.cpp_extension import _get_build_directory - - lock_path = os.path.join(_get_build_directory(name, False), "lock") - start = time.time() - while os.path.exists(lock_path): - if time.time() - start > timeout: - raise RuntimeError( - f"Timed out after {timeout} seconds waiting for another process " - f"to finish building the OpenEquivariance extension '{name}'. " - f"If no other build is running, a previous build was likely " - f"killed partway; delete {lock_path} and retry." - ) - time.sleep(1) - - -def load_jit_extension(): - global \ - BUILT_EXTENSION, \ - BUILT_EXTENSION_ERROR, \ - LINKED_LIBPYTHON, \ - LINKED_LIBPYTHON_ERROR, \ - extension_module - - # Locate libpython (required for AOTI) - try: - python_lib_dir = sysconfig.get_config_var("LIBDIR") - major, minor = sys.version_info.major, sys.version_info.minor - python_lib_name = f"python{major}.{minor}" - - libpython_so = os.path.join(python_lib_dir, f"lib{python_lib_name}.so") - libpython_a = os.path.join(python_lib_dir, f"lib{python_lib_name}.a") - if not (os.path.exists(libpython_so) or os.path.exists(libpython_a)): - raise FileNotFoundError( - f"libpython not found, tried {libpython_so} and {libpython_a}" - ) +SOURCE_BUILD_HINT = ( + "To build the extension for your setup from source, run:\n" + " pip install --no-binary openequivariance openequivariance" +) - LINKED_LIBPYTHON = True - except Exception as e: - LINKED_LIBPYTHON_ERROR = f"Error linking libpython:\n{e}\nSysconfig variables:\n{sysconfig.get_config_vars()}" - try: - from torch.utils.cpp_extension import library_paths, include_paths +class Extension: + def __init__(self): + self.built_extension = False + self.built_extension_error = None + self.ext_module_python = None + self.aoti_module_path = None - extra_cflags = ["-O3"] - torch_sources = ["libtorch_tp_jit.cpp", "json11/json11.cpp"] - include_dirs, extra_link_args = (["backend"], ["-Wl,--no-as-needed"]) - extra_include_dirs = [] +class JITExtension(Extension): + """ + Only SYCL uses PyTorch's JIT compilation flow; CUDA and HIP now use + precompiled extensions shipped with the wheel. + """ + @staticmethod + @contextlib.contextmanager + def maybe_set_max_jobs_env(limit=8): + if "MAX_JOBS" in os.environ: + yield + return + os.environ["MAX_JOBS"] = str(min(limit, os.cpu_count() or 1)) try: - import pybind11 + yield + finally: + os.environ.pop("MAX_JOBS", None) + + @staticmethod + def wait_for_torch_build_lock(name, timeout=300): + from torch.utils.cpp_extension import _get_build_directory + + lock_path = os.path.join(_get_build_directory(name, False), "lock") + start = time.time() + while os.path.exists(lock_path): + if time.time() - start > timeout: + raise RuntimeError( + f"Timed out after {timeout} seconds waiting for another process " + f"to finish building the OpenEquivariance extension '{name}'. " + f"If no other build is running, a previous build was likely " + f"killed partway; delete {lock_path} and retry." + ) + time.sleep(1) + + def __init__(self): + super().__init__() + python_lib_dir = sysconfig.get_config_var("LIBDIR") or "" + python_lib_name = f"python{sys.version_info.major}.{sys.version_info.minor}" + found_libpython = any( + os.path.exists(os.path.join(python_lib_dir, f"lib{python_lib_name}{ext}")) + for ext in (".so", ".a") + ) - extra_include_dirs.append(pybind11.get_include()) - except Exception as e: - BUILT_EXTENSION_ERROR = ( - "Could not locate pybind11 include path required for JIT " - f"OpenEquivariance extension compilation: {e}" + missing = [ + name + for name in ("setuptools", "ninja") + if importlib.util.find_spec(name) is None + ] + if missing: + self.built_extension_error = ( + f"The SYCL backend compiles its extension at first import, which " + f"needs {', '.join(missing)}. Install them with:\n" + " pip install openequivariance[sycl]" ) return - if LINKED_LIBPYTHON: - extra_link_args.pop() - extra_link_args.extend( - [ - f"-Wl,--no-as-needed,-rpath,{python_lib_dir}", - f"-L{python_lib_dir}", - f"-l{python_lib_name}", - ], - ) - if BACKEND == "cuda": - extra_link_args.extend(["-lcuda", "-lcudart", "-lnvrtc", "-ltorch_cuda"]) - - try: - torch_libs, cuda_libs = library_paths("cuda") - extra_link_args.append("-Wl,-rpath," + torch_libs) - extra_link_args.append("-L" + cuda_libs) - if os.path.exists(cuda_libs + "/stubs"): - extra_link_args.append("-L" + cuda_libs + "/stubs") - except Exception as e: - getLogger().info(str(e)) - - extra_cflags.append("-DCUDA_BACKEND") - elif BACKEND == "hip": - torch_libs = library_paths("cuda")[0] - extra_link_args.append("-Wl,-rpath," + torch_libs) - extra_link_args.extend("-L" + path for path in library_paths("cuda")) - extra_link_args.extend(["-ltorch_hip", "-lhiprtc"]) - extra_cflags.append("-DHIP_BACKEND") - elif BACKEND == "sycl": + try: + from torch.utils.cpp_extension import library_paths, include_paths + + extra_cflags = ["-O3"] + torch_sources = ["libtorch_tp_jit.cpp", "json11/json11.cpp"] + + include_dirs, extra_link_args = (["backend"], ["-Wl,--no-as-needed"]) + + if found_libpython: + extra_link_args.pop() + extra_link_args.extend( + [ + f"-Wl,--no-as-needed,-rpath,{python_lib_dir}", + f"-L{python_lib_dir}", + f"-l{python_lib_name}", + ], + ) import shutil cxx = os.environ.get("CXX", "") if "icpx" not in os.path.basename(cxx): if shutil.which("icpx") is None: - BUILT_EXTENSION_ERROR = ( + self.built_extension_error = ( "The SYCL backend requires the oneAPI DPC++ compiler. " "Put 'icpx' on your PATH or set CXX to it." ) @@ -173,107 +124,111 @@ def load_jit_extension(): extra_link_args.append("-Wl,-rpath," + lib_dir) extra_link_args.append("-L" + lib_dir) - torch_sources = [oeq_root + "/extension/" + src for src in torch_sources] - include_dirs = ( - [oeq_root + "/extension/" + d for d in include_dirs] - + extra_include_dirs - + include_paths("xpu" if BACKEND == "sycl" else "cuda") - ) - - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - - try: - _wait_for_torch_build_lock("libtorch_tp_jit") - with _maybe_set_max_jobs_env(): - extension_module = torch.utils.cpp_extension.load( - "libtorch_tp_jit", - torch_sources, - extra_cflags=extra_cflags, - extra_include_paths=include_dirs, - extra_ldflags=extra_link_args, - ) - torch.ops.load_library(extension_module.__file__) - BUILT_EXTENSION = True - except Exception as e: - # If compiling torch fails (e.g. low gcc version), we should fall back to the - # version that takes integer pointers as args (but is untraceable to PyTorch JIT / export). - BUILT_EXTENSION_ERROR = e - except Exception as e: - BUILT_EXTENSION_ERROR = f"Error JIT-compiling OpenEquivariance Extension: {e}" - - -def load_precompiled_extension(): - global BUILT_EXTENSION, BUILT_EXTENSION_ERROR, LINKED_LIBPYTHON, extension_module - LINKED_LIBPYTHON = ( - True # Doesn't actually use libpython, just set this as true anyway - ) - try: - if BACKEND == "cuda": - import openequivariance._torch.extlib.oeq_stable_cuda as extension_module - elif BACKEND == "hip": - import openequivariance._torch.extlib.oeq_stable_hip as extension_module - elif BACKEND == "sycl": - import openequivariance._torch.extlib.oeq_stable_sycl as extension_module - - torch.ops.load_library(extension_module.__file__) - BUILT_EXTENSION = True - except Exception as e: - BUILT_EXTENSION_ERROR = ( - f"Error loading precompiled OpenEquivariance Extension: {e}" - ) - + torch_sources = [oeq_root + "/extension/" + src for src in torch_sources] + include_dirs = [ + oeq_root + "/extension/" + d for d in include_dirs + ] + include_paths("xpu") + + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + + try: + self.wait_for_torch_build_lock("libtorch_tp_jit") + with self.maybe_set_max_jobs_env(): + self.ext_module_python = torch.utils.cpp_extension.load( + "libtorch_tp_jit", + torch_sources, + extra_cflags=extra_cflags, + extra_include_paths=include_dirs, + extra_ldflags=extra_link_args, + ) + torch.ops.load_library(self.ext_module_python.__file__) + self.built_extension = True + self.aoti_module_path = self.ext_module_python.__file__ + + except Exception as e: + # If compiling torch fails (e.g. low gcc version), we should fall back to the + # version that takes integer pointers as args (but is untraceable to PyTorch JIT / export). + self.built_extension_error = e + except Exception as e: + self.built_extension_error = ( + f"Error JIT-compiling OpenEquivariance Extension: {e}" + ) -USE_PRECOMPILED_EXTENSION = True -WARNING_MESSAGE = "" -if os.getenv("OEQ_JIT_EXTENSION", "0") == "1": - WARNING_MESSAGE += "Environment variable OEQ_JIT_EXTENSION=1 is set.\n" - USE_PRECOMPILED_EXTENSION = False +class PrecompiledExtension(Extension): + def __init__(self, backend): + super().__init__() + if backend == "cuda": + ext_module_py_name = f"oeq_stable_cuda{torch.version.cuda.split('.')[0]}" + else: + ext_module_py_name = f"oeq_stable_{backend}" -if Version(torch.__version__) <= Version("2.9.9"): - WARNING_MESSAGE += f"PyTorch version {torch.__version__} is < 2.10, minimum required for precompiled extension. Please upgrade to 2.10.\n" - USE_PRECOMPILED_EXTENSION = False + full_name = f"openequivariance._torch.extlib.{ext_module_py_name}" + spec = importlib.util.find_spec(full_name) + aoti_module_path = spec and os.path.join( + os.path.dirname(spec.origin), f"lib{ext_module_py_name}_aoti.so" + ) + if spec is None or not os.path.exists(aoti_module_path): + self.built_extension_error = ( + f"No precompiled OpenEquivariance extension {ext_module_py_name} " + f"for PyTorch {torch.__version__}. {SOURCE_BUILD_HINT}" + ) + return -if BACKEND == "hip": - WARNING_MESSAGE += "HIP does not support precompiled extension yet.\n" - USE_PRECOMPILED_EXTENSION = False + try: + self.ext_module_python = importlib.import_module(full_name) + torch.ops.load_library(self.ext_module_python.__file__) + self.built_extension = True + self.aoti_module_path = aoti_module_path + except Exception as e: + self.built_extension_error = ( + f"Error loading precompiled OpenEquivariance extension " + f"{ext_module_py_name}: {e}\n{SOURCE_BUILD_HINT}" + ) -if BACKEND == "sycl": - WARNING_MESSAGE += "SYCL does not support precompiled extension yet.\n" - USE_PRECOMPILED_EXTENSION = False -AOTI_SO_NAME = f"liboeq_stable_{BACKEND}_aoti.so" +def _detect_backend(): + """HIP builds report a ``torch.version.cuda`` of None, so HIP is tested first.""" + if torch.version.hip: + return "hip" + if torch.version.cuda: + return "cuda" + if getattr(torch.version, "xpu", None): + return "sycl" + raise RuntimeError( + "Only the CUDA, HIP and XPU (SYCL) backends are supported. " + "No supported accelerator was detected in this PyTorch build." + ) -if not os.path.exists(os.path.join(os.path.dirname(__file__), AOTI_SO_NAME)): - WARNING_MESSAGE += "Precompiled extension shared object not found.\n" - USE_PRECOMPILED_EXTENSION = False +BACKEND = _detect_backend() +IS_HIP = BACKEND == "hip" +IS_SYCL = BACKEND == "sycl" +DEVICE_TYPE = "xpu" if IS_SYCL else "cuda" -if USE_PRECOMPILED_EXTENSION: - load_precompiled_extension() +extension = None +if BACKEND == "sycl": + extension = JITExtension() else: - WARNING_MESSAGE += "For these reasons, falling back to JIT compilation of OpenEquivariance extension. If another process holds the build lock, this waits up to 5 minutes, then raises an error naming the lock file to delete.\n" - warnings.warn(WARNING_MESSAGE, stacklevel=3) - load_jit_extension() - - -def torch_ext_so_path(): - if not USE_PRECOMPILED_EXTENSION: - return extension_module.__file__ - else: - dirname = os.path.dirname(extension_module.__file__) - return os.path.join(dirname, AOTI_SO_NAME) + extension = PrecompiledExtension(BACKEND) +BUILT_EXTENSION = extension.built_extension +BUILT_EXTENSION_ERROR = extension.built_extension_error -sys.modules["oeq_utilities"] = extension_module - if BUILT_EXTENSION: + sys.modules["oeq_utilities"] = extension.ext_module_python from oeq_utilities import ( DeviceProp, GPUTimer, ) + + def torch_ext_so_path(): + return extension.aoti_module_path + else: + if BUILT_EXTENSION_ERROR is not None: + warnings.warn(str(BUILT_EXTENSION_ERROR), stacklevel=3) def _raise_import_error_helper(import_target: str): if not BUILT_EXTENSION: @@ -286,3 +241,6 @@ def DeviceProp(*args, **kwargs): def GPUTimer(*args, **kwargs): _raise_import_error_helper("GPUTimer") + + def torch_ext_so_path(): + _raise_import_error_helper("torch_ext_so_path") diff --git a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp index 980a55d8..7e99c0ce 100644 --- a/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp +++ b/openequivariance/openequivariance/extension/libtorch_tp_jit_stable.cpp @@ -140,34 +140,22 @@ bool tensor_is_on_gpu(const Tensor &tensor) { } // namespace oeq -#ifdef CUDA_BACKEND - #define EXTENSION_NAME oeq_stable_cuda -#endif -#ifdef HIP_BACKEND - #define EXTENSION_NAME oeq_stable_hip +// CMake passes EXTENSION_NAME to match the target name (e.g. oeq_stable_cuda13) +#ifndef EXTENSION_NAME + #ifdef CUDA_BACKEND + #define EXTENSION_NAME oeq_stable_cuda + #endif + #ifdef HIP_BACKEND + #define EXTENSION_NAME oeq_stable_hip + #endif + #ifdef SYCL_BACKEND + #define EXTENSION_NAME oeq_stable_sycl + #endif #endif -#ifdef SYCL_BACKEND - #define EXTENSION_NAME oeq_stable_sycl -#endif #ifdef INCLUDE_NB_EXTENSION - #include "nanobind/nanobind.h" - #include "nanobind/stl/string.h" - namespace nb = nanobind; + #include "nb_utilities.hpp" NB_MODULE(EXTENSION_NAME, m) { - nb::class_(m, "DeviceProp") - .def(nb::init()) - .def_ro("name", &DeviceProp::name) - .def_ro("warpsize", &DeviceProp::warpsize) - .def_ro("major", &DeviceProp::major) - .def_ro("minor", &DeviceProp::minor) - .def_ro("multiprocessorCount", &DeviceProp::multiprocessorCount) - .def_ro("maxSharedMemPerBlock", &DeviceProp::maxSharedMemPerBlock); - - nb::class_(m, "GPUTimer") - .def(nb::init<>()) - .def("start", &GPUTimer::start) - .def("stop_clock_get_elapsed", &GPUTimer::stop_clock_get_elapsed) - .def("clear_L2_cache", &GPUTimer::clear_L2_cache); + oeq::bind_utilities(m); } #endif diff --git a/openequivariance/openequivariance/extension/nb_utilities.hpp b/openequivariance/openequivariance/extension/nb_utilities.hpp new file mode 100644 index 00000000..e5c3d5cf --- /dev/null +++ b/openequivariance/openequivariance/extension/nb_utilities.hpp @@ -0,0 +1,35 @@ +#pragma once + +#include "nanobind/nanobind.h" +#include "nanobind/stl/string.h" + +namespace oeq { + +inline void bind_utilities(nanobind::module_ &m) { + namespace nb = nanobind; + + if (nb::handle existing = nb::type()) { + m.attr("DeviceProp") = existing; + } else { + nb::class_(m, "DeviceProp") + .def(nb::init()) + .def_ro("name", &DeviceProp::name) + .def_ro("warpsize", &DeviceProp::warpsize) + .def_ro("major", &DeviceProp::major) + .def_ro("minor", &DeviceProp::minor) + .def_ro("multiprocessorCount", &DeviceProp::multiprocessorCount) + .def_ro("maxSharedMemPerBlock", &DeviceProp::maxSharedMemPerBlock); + } + + if (nb::handle existing = nb::type()) { + m.attr("GPUTimer") = existing; + } else { + nb::class_(m, "GPUTimer") + .def(nb::init<>()) + .def("start", &GPUTimer::start) + .def("stop_clock_get_elapsed", &GPUTimer::stop_clock_get_elapsed) + .def("clear_L2_cache", &GPUTimer::clear_L2_cache); + } +} + +} // namespace oeq diff --git a/openequivariance/pyproject.toml b/openequivariance/pyproject.toml index 4e573d2e..4fee2c0a 100644 --- a/openequivariance/pyproject.toml +++ b/openequivariance/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools", "setuptools-scm", "scikit-build-core", "nanobind"] +requires = ["setuptools", "setuptools-scm", "scikit-build-core", "nanobind>=3.0"] build-backend = "scikit_build_core.build" [project] @@ -12,13 +12,11 @@ authors = [ { name="James Demmel" } ] description = "A fast GPU JIT kernel generator for the Clebsch-Gordon Tensor Product" -requires-python = ">=3.10" +requires-python = ">=3.11" dependencies = [ - "setuptools", - "ninja", "jinja2", "numpy", - "pybind11" + "nanobind-backend>=1.0" ] readme = "README.md" @@ -27,7 +25,6 @@ license-files = ["LICENSE"] classifiers = [ "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -60,6 +57,11 @@ dev = [ "cmake" ] +sycl = [ + "setuptools", + "ninja" +] + jax = [ "scikit-build-core", "setuptools-scm", @@ -77,5 +79,41 @@ addopts = [ [tool.ruff] lint.ignore = ["E741"] +[tool.scikit-build] +wheel.py-api = "cp311" +build-dir = "build/{wheel_tag}" +editable.rebuild = true + [tool.scikit-build.cmake.define] AARCH64 = {env="AARCH64", default="0"} +LIBTORCH_URL_X86_64 = "https://download.pytorch.org/libtorch/cpu/libtorch-shared-with-deps-2.10.0%2Bcpu.zip" +LIBTORCH_SHA256_X86_64 = "c5bf8efda9224a2d971b19d1ef6cf3ba6fee8ab53e69c49427db003d1d300496" +LIBTORCH_URL_AARCH64 = "https://github.com/PASSIONLab/libtorch_aarch64/raw/742ffd07f47a31a94f6b6aae2cbe129e6b667d94/pkg-libtorch.tar.zst" +LIBTORCH_SHA256_AARCH64 = "9fdebdb56766317c71eb4c9f1b750baef4ebc4e6ba85b427f701fdc20844f909" + +[tool.cibuildwheel] +build = "cp311-*" +skip = "*-musllinux*" +manylinux-x86_64-image = "manylinux_2_28" +manylinux-aarch64-image = "manylinux_2_28" +before-all = [ + "/opt/python/cp312-cp312/bin/python {project}/.github/release.py install-cuda-toolkits $OEQ_CUDA_ROOTS", + "/opt/python/cp312-cp312/bin/python {project}/.github/release.py install-rocm-sdk {package}/pyproject.toml /opt/rocm-sdk-core", +] +build-verbosity = 1 + +[tool.cibuildwheel.linux.environment] +OEQ_CUDA_ROOTS = "/usr/local/cuda-12.6 /usr/local/cuda-13.0" +ROCM_PATH = "/opt/rocm-sdk-core/_rocm_sdk_core" + +[tool.cibuildwheel.linux] +repair-wheel-command = """auditwheel repair -w {dest_dir} {wheel} \ + --exclude libtorch_cpu.so --exclude libc10.so \ + --exclude libtorch_cuda.so --exclude libtorch_hip.so \ + --exclude libcuda.so.1 --exclude 'libcudart.so.*' --exclude 'libnvrtc.so.*' \ + --exclude 'libamdhip64.so.*' --exclude 'libhiprtc.so.*'""" + +[tool.openequivariance.rocm-sdk-core] +index-url = "https://download.pytorch.org/whl/rocm10.0" +version = "10.0.0" +sha256 = "3bf4a72d11aa2a4ee1e90572c73630f937d38c1b7af4dda45b483b54655138a7" diff --git a/openequivariance_extjax/src/jax_ffi_registration.cpp b/openequivariance_extjax/src/jax_ffi_registration.cpp index 86307f95..83d8d12f 100644 --- a/openequivariance_extjax/src/jax_ffi_registration.cpp +++ b/openequivariance_extjax/src/jax_ffi_registration.cpp @@ -10,6 +10,8 @@ #include "backend/backend_hip.hpp" #endif +#include "nb_utilities.hpp" + namespace nb = nanobind; namespace { @@ -47,18 +49,5 @@ NB_MODULE(openequivariance_extjax, m) { }); m.def("is_hip", &IsHip); - nb::class_(m, "DeviceProp") - .def(nb::init()) - .def_ro("name", &DeviceProp::name) - .def_ro("warpsize", &DeviceProp::warpsize) - .def_ro("major", &DeviceProp::major) - .def_ro("minor", &DeviceProp::minor) - .def_ro("multiprocessorCount", &DeviceProp::multiprocessorCount) - .def_ro("maxSharedMemPerBlock", &DeviceProp::maxSharedMemPerBlock); - - nb::class_(m, "GPUTimer") - .def(nb::init<>()) - .def("start", &GPUTimer::start) - .def("stop_clock_get_elapsed", &GPUTimer::stop_clock_get_elapsed) - .def("clear_L2_cache", &GPUTimer::clear_L2_cache); + oeq::bind_utilities(m); } diff --git a/tests/export_test.py b/tests/export_test.py index 1b77df7a..f84a5310 100644 --- a/tests/export_test.py +++ b/tests/export_test.py @@ -144,7 +144,6 @@ def test_aoti(tp_and_inputs): def test_aoti_cpp_inference(problem_and_irreps): - assert oeq.LINKED_LIBPYTHON, oeq.LINKED_LIBPYTHON_ERROR problem, X_ir, Y_ir, _ = problem_and_irreps cmake_prefix_path = torch.utils.cmake_prefix_path torch_ext_so_path = oeq.torch_ext_so_path()