diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 00000000..877f45f1 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,191 @@ +name: Build wheels and publish + +on: + workflow_dispatch: + workflow_call: + +jobs: + build: + name: Build wheels on ${{ matrix.os }}-${{ matrix.accelerator }} + runs-on: ${{ matrix.os }} + strategy: + fail-fast: false + matrix: + os: [ubuntu-latest, ubuntu-24.04-arm, windows-2022, macos-latest] + accelerator: [cpu, cu118, cu126] #, cu128] + exclude: + - os: ubuntu-24.04-arm + accelerator: cu118 + - os: ubuntu-24.04-arm + accelerator: cu126 + # - os: ubuntu-24.04-arm + # accelerator: cu128 + - os: macos-latest + accelerator: cu118 + - os: macos-latest + accelerator: cu126 + # - os: macos-latest + # accelerator: cu128 + + steps: + - name: Free space of Github Runner (otherwise it will fail by running out of disk) + if: matrix.os == 'ubuntu-latest' + run: | + sudo rm -rf /usr/share/dotnet + sudo rm -rf /opt/ghc + sudo rm -rf "/usr/local/share/boost" + sudo rm -rf "/usr/local/.ghcup" + sudo rm -rf "/usr/local/julia1.9.2" + sudo rm -rf "/usr/local/lib/android" + sudo rm -rf "$AGENT_TOOLSDIRECTORY" + + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.13" + + - name: Install cibuildwheel + run: python -m pip install cibuildwheel==2.23.0 + + - name: Activate MSVC + uses: ilammy/msvc-dev-cmd@v1 + with: + toolset: 14.29 + if: matrix.os == 'windows-2022' + + - name: "Install dependencies with choco" + uses: crazy-max/ghaction-chocolatey@v3 + if: matrix.os == 'windows-2022' + with: + args: install jom + + - name: Build wheels + if: matrix.os != 'windows-2022' + shell: bash -l {0} + run: python -m cibuildwheel python/ --output-dir wheelhouse + env: + ACCELERATOR: ${{ matrix.accelerator }} + + - name: Build wheels + if: matrix.os == 'windows-2022' + shell: cmd # Use cmd on Windows to avoid bash environment taking priority over MSVC variables + run: python -m cibuildwheel python/ --output-dir wheelhouse + env: + DISTUTILS_USE_SDK: "1" # Windows requires this to use vc for building + ACCELERATOR: ${{ matrix.accelerator }} + + - uses: actions/upload-artifact@v4 + with: + name: ${{ matrix.accelerator }}-cibw-wheels-${{ matrix.os }}-${{ strategy.job-index }} + path: ./wheelhouse/*.whl + + publish-to-public-pypi: + name: >- + Publish Python 🐍 distribution 📦 to PyPI + needs: + - build + runs-on: ubuntu-latest + environment: + name: pypi + permissions: + id-token: write # IMPORTANT: mandatory for trusted publishing + strategy: + fail-fast: false + matrix: + accelerator: [cpu, cu118, cu126] #, hip, cu124, cu126, cu128] + + steps: + - name: Download all the dists + uses: actions/download-artifact@v4 + with: + pattern: "${{ matrix.accelerator }}-cibw-wheels*" + path: dist/ + merge-multiple: true + + - name: Publish distribution 📦 to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 + with: + password: ${{ secrets.PUBLIC_PYPI_API_TOKEN }} + skip-existing: true + + # publish-to-accelera-pypi: + # name: >- + # Publish Python 🐍 distribution 📦 to Acellera PyPI + # needs: + # - build + # runs-on: ubuntu-latest + # permissions: # Needed for GCP authentication + # contents: "read" + # id-token: "write" + # strategy: + # fail-fast: false + # matrix: + # accelerator: [cpu, cu118, cu126] #, cu128] + + # steps: + # - uses: actions/checkout@v4 # Needed for GCP authentication for some reason + + # - name: Set up Cloud SDK + # uses: google-github-actions/auth@v2 + # with: + # workload_identity_provider: ${{ secrets.GCP_WORKLOAD_IDENTITY_PROVIDER }} + # service_account: ${{ secrets.GCP_PYPI_SERVICE_ACCOUNT }} + + # - name: Download all the dists + # uses: actions/download-artifact@v4 + # with: + # pattern: "${{ matrix.accelerator }}-cibw-wheels*" + # path: dist/ + # merge-multiple: true + + # - name: Publish distribution 📦 to Acellera PyPI + # run: | + # pip install build twine keyring keyrings.google-artifactregistry-auth + # pip install -U packaging + # twine upload --repository-url https://us-central1-python.pkg.dev/pypi-packages-455608/${{ matrix.accelerator }} dist/* --verbose --skip-existing + + github-release: + name: >- + Sign the Python 🐍 distribution 📦 with Sigstore + and upload them to GitHub Release + needs: + - build + runs-on: ubuntu-latest + + permissions: + contents: write # IMPORTANT: mandatory for making GitHub Releases + id-token: write # IMPORTANT: mandatory for sigstore + + steps: + - name: Download all the dists + uses: actions/download-artifact@v4 + with: + path: dist/ + merge-multiple: true + + - name: Sign the dists with Sigstore + uses: sigstore/gh-action-sigstore-python@v3.0.0 + with: + inputs: >- + ./dist/*.whl + + - name: Create GitHub Release + env: + GITHUB_TOKEN: ${{ github.token }} + run: >- + gh release create + "$GITHUB_REF_NAME" + --repo "$GITHUB_REPOSITORY" + --notes "" + + - name: Upload artifact signatures to GitHub Release + env: + GITHUB_TOKEN: ${{ github.token }} + # Upload to GitHub Release using the `gh` CLI. + # `dist/` contains the built packages, and the + # sigstore-produced signatures and certificates. + run: >- + gh release upload + "$GITHUB_REF_NAME" dist/** + --repo "$GITHUB_REPOSITORY" diff --git a/.github/workflows/workflow.yml b/.github/workflows/workflow.yml new file mode 100644 index 00000000..45eab0be --- /dev/null +++ b/.github/workflows/workflow.yml @@ -0,0 +1,9 @@ +name: Whole + +on: [push] + +jobs: + publish: + if: startsWith(github.event.ref, 'refs/tags/v') + uses: ./.github/workflows/publish.yml + secrets: inherit diff --git a/CMakeLists.txt b/CMakeLists.txt index 75095427..2442b89e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -92,6 +92,7 @@ SET_TARGET_PROPERTIES(${SHARED_NN_TARGET} LINK_FLAGS "${EXTRA_COMPILE_FLAGS}") TARGET_LINK_LIBRARIES(${SHARED_NN_TARGET} OpenMM) TARGET_LINK_LIBRARIES(${SHARED_NN_TARGET} "${TORCH_LIBRARIES}") +SET_TARGET_PROPERTIES(${SHARED_NN_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/../../torch/lib;$ORIGIN/../../nvidia/cuda_nvrtc/lib") INSTALL_TARGETS(/lib RUNTIME_DIRECTORY /lib ${SHARED_NN_TARGET}) # install headers diff --git a/platforms/cuda/CMakeLists.txt b/platforms/cuda/CMakeLists.txt index 0eb9e831..c24451e1 100644 --- a/platforms/cuda/CMakeLists.txt +++ b/platforms/cuda/CMakeLists.txt @@ -76,6 +76,7 @@ IF (APPLE) SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES LINK_FLAGS "-F/Library/Frameworks -framework CUDA ${EXTRA_COMPILE_FLAGS}") ENDIF (APPLE) +SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;$ORIGIN/../../../nvidia/cuda_runtime/lib/;$ORIGIN/../../../torch/lib;") INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins) SUBDIRS (tests) diff --git a/platforms/hip/CMakeLists.txt b/platforms/hip/CMakeLists.txt index fbe97ee9..b370af00 100644 --- a/platforms/hip/CMakeLists.txt +++ b/platforms/hip/CMakeLists.txt @@ -60,6 +60,7 @@ SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}" LINK_FLAGS "${EXTRA_COMPILE_FLAGS}") +SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;$ORIGIN/../../../torch/lib;") INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins) SUBDIRS (tests) diff --git a/platforms/opencl/CMakeLists.txt b/platforms/opencl/CMakeLists.txt index 506ae5c0..2eba1b4c 100644 --- a/platforms/opencl/CMakeLists.txt +++ b/platforms/opencl/CMakeLists.txt @@ -60,6 +60,7 @@ SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}" LINK_FLAGS "${EXTRA_COMPILE_FLAGS}") +SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;$ORIGIN/../../../torch/lib;") INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins) SUBDIRS (tests) diff --git a/platforms/reference/CMakeLists.txt b/platforms/reference/CMakeLists.txt index 124c82fb..c13f564b 100644 --- a/platforms/reference/CMakeLists.txt +++ b/platforms/reference/CMakeLists.txt @@ -46,6 +46,7 @@ SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES COMPILE_FLAGS "-DOPENMM_BUILDING_SHARED_LIBRARY ${EXTRA_COMPILE_FLAGS}" LINK_FLAGS "${EXTRA_COMPILE_FLAGS}") +SET_TARGET_PROPERTIES(${SHARED_TARGET} PROPERTIES INSTALL_RPATH "$ORIGIN;$ORIGIN/..;$ORIGIN/../../../torch/lib;") INSTALL(TARGETS ${SHARED_TARGET} DESTINATION ${CMAKE_INSTALL_PREFIX}/lib/plugins) SUBDIRS (tests) diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index d7675cdc..77c88a56 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -24,7 +24,7 @@ add_custom_target(PythonInstall DEPENDS "${WRAP_FILE}" "${CMAKE_CURRENT_SOURCE_ set(NN_PLUGIN_HEADER_DIR "${CMAKE_SOURCE_DIR}/openmmapi/include") set(NN_PLUGIN_LIBRARY_DIR "${CMAKE_BINARY_DIR}") configure_file(${CMAKE_CURRENT_SOURCE_DIR}/setup.py ${CMAKE_CURRENT_BINARY_DIR}/setup.py) -add_custom_command(TARGET PythonInstall - COMMAND "${PYTHON_EXECUTABLE}" -m pip install . - WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" -) +# add_custom_command(TARGET PythonInstall +# COMMAND "${PYTHON_EXECUTABLE}" -m pip install . +# WORKING_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}" +# ) diff --git a/python/MANIFEST.in b/python/MANIFEST.in new file mode 100644 index 00000000..8a69508d --- /dev/null +++ b/python/MANIFEST.in @@ -0,0 +1,4 @@ +graft openmm/ +graft openmmtorch/ +prune tests +prune cibuildwheel_support \ No newline at end of file diff --git a/python/cibuildwheel_support/before_all_linux.sh b/python/cibuildwheel_support/before_all_linux.sh new file mode 100755 index 00000000..4de33b31 --- /dev/null +++ b/python/cibuildwheel_support/before_all_linux.sh @@ -0,0 +1,123 @@ +#! /bin/bash + +set -e +set -x + +# Install dependencies with yum +dnf install -y zip opencl-headers ocl-icd tree + +# Check if we are running on aarch64 +if [ "$(uname -m)" == "aarch64" ]; then + dnf install -y libgfortran + ln -s /usr/lib64/libgfortran.so.5.0.0 /usr/lib64/libgfortran-0b50f350.so.5.0.0 +fi + +# Configure pip to use PyTorch extra-index-url for CPU +mkdir -p $HOME/.config/pip +echo "[global] +extra-index-url = https://download.pytorch.org/whl/cpu" > $HOME/.config/pip/pip.conf + +if [ "$ACCELERATOR" == "cu118" ]; then + # Install CUDA 11.8 + dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo + + dnf install --setopt=obsoletes=0 -y \ + cuda-compiler-11-8-11.8.0-1 \ + cuda-libraries-11-8-11.8.0-1 \ + cuda-libraries-devel-11-8-11.8.0-1 \ + cuda-toolkit-11-8-11.8.0-1 \ + gcc-toolset-11 + + ln -s cuda-11.8 /usr/local/cuda + ln -s /opt/rh/gcc-toolset-11/root/usr/bin/gcc /usr/local/cuda/bin/gcc + ln -s /opt/rh/gcc-toolset-11/root/usr/bin/g++ /usr/local/cuda/bin/g++ + ln -s /usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so /usr/lib/libcuda.so.1 + + # Configure pip to use PyTorch extra-index-url for CUDA 11.8 + mkdir -p $HOME/.config/pip + echo "[global] +extra-index-url = https://download.pytorch.org/whl/cu118" > $HOME/.config/pip/pip.conf + + pip install openmm-unofficial-cu11 torch==2.7.1 +elif [ "$ACCELERATOR" == "cu126" ]; then + # Install CUDA 12.6 + dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo + + dnf install --setopt=obsoletes=0 -y \ + cuda-compiler-12-6-12.6.3-1\ + cuda-libraries-12-6-12.6.3-1 \ + cuda-libraries-devel-12-6-12.6.3-1 \ + cuda-toolkit-12-6-12.6.3-1 \ + gcc-toolset-13 + + ln -s cuda-12.6 /usr/local/cuda + ln -s /opt/rh/gcc-toolset-13/root/usr/bin/gcc /usr/local/cuda/bin/gcc + ln -s /opt/rh/gcc-toolset-13/root/usr/bin/g++ /usr/local/cuda/bin/g++ + ln -s /usr/local/cuda/targets/x86_64-linux/lib/stubs/libcuda.so /usr/lib/libcuda.so.1 + + # Configure pip to use PyTorch extra-index-url for CUDA 12.6 + mkdir -p $HOME/.config/pip + echo "[global] +extra-index-url = https://download.pytorch.org/whl/cu126" > $HOME/.config/pip/pip.conf + + pip install openmm-unofficial-cu12 torch==2.7.1 +elif [ "$ACCELERATOR" == "hip" ]; then + # Install HIP 6.2 + dnf install -y https://repo.radeon.com/amdgpu-install/6.2.2/el/8.10/amdgpu-install-6.2.60202-1.el8.noarch.rpm + dnf install -y rocm-device-libs hip-devel hip-runtime-amd hipcc + pip install openmm-unofficial-cpu torch==2.7.1 +else + pip install openmm-unofficial-cpu torch==2.7.1 +fi + +################# +SITE_PACKAGES=$(python -c 'import site; print(site.getsitepackages()[0])') + +CMAKE_FLAGS="" +if [ "$ACCELERATOR" == "cu118" ] || [ "$ACCELERATOR" == "cu126" ]; then + ARCH_LIST=$(python -c "import torch; print(';'.join([f'{y[:-1]}.{y[-1]}' for y in [x[3:] for x in torch._C._cuda_getArchFlags().split() if x.startswith('sm_')]]))") + # CMakeLists.txt seems to ignore the CMAKE_CUDA_ARCHITECTURES variable, instead, it is overwritten by TORCH_CUDA_ARCH_LIST + ARCH_LIST_FMT=$(python -c "import torch; print(';'.join([y for y in [x[3:] for x in torch._C._cuda_getArchFlags().split() if x.startswith('sm_')]]))") + export CUDA_HOME="/usr/local/cuda" + CMAKE_FLAGS=" -DTORCH_CUDA_ARCH_LIST=${ARCH_LIST}" + CMAKE_FLAGS+=" -DCMAKE_CUDA_ARCHITECTURES=${ARCH_LIST_FMT}" + CMAKE_FLAGS+=" -DCUDA_TOOLKIT_ROOT_DIR=${CUDA_HOME}" + CMAKE_FLAGS+=" -DCMAKE_CUDA_COMPILER=${CUDA_HOME}/bin/nvcc" +fi + +if [[ "$OSTYPE" == "linux-gnu"* ]]; then + CMAKE_FLAGS+=" -DOPENCL_INCLUDE_DIR=/usr/include/CL" + CMAKE_FLAGS+=" -DOPENCL_LIBRARY=/usr/lib64/libOpenCL.so.1" +fi + +# Configure build with Cmake +mkdir -p build +mkdir -p install +cd build + +echo $CMAKE_FLAGS +export LD_LIBRARY_PATH=/usr/lib/:/usr/local/cuda/targets/x86_64-linux/lib/:$LD_LIBRARY_PATH + +cmake .. \ + -DCMAKE_INSTALL_PREFIX=../install \ + -DCMAKE_BUILD_TYPE=Release \ + -DOPENMM_DIR=${SITE_PACKAGES}/openmm \ + -DPYTORCH_DIR=${SITE_PACKAGES}/torch \ + -DTorch_DIR=${SITE_PACKAGES}/torch/share/cmake/Torch \ + -DNN_BUILD_OPENCL_LIB=ON \ + ${CMAKE_FLAGS} + +# Build OpenMMTorch +make -j4 install +make -j4 PythonInstall + +cd .. + +cp build/python/setup.py python/ +cp build/python/openmmtorch.py python/openmmtorch/ +cp build/python/TorchPluginWrapper.cpp python/openmmtorch/ +cp -r install/include python/openmmtorch/ + +# Copy the libraries of openmm +mkdir -p python/openmm/ +cp -r install/lib python/openmm/lib \ No newline at end of file diff --git a/python/cibuildwheel_support/before_all_osx.sh b/python/cibuildwheel_support/before_all_osx.sh new file mode 100755 index 00000000..7812424f --- /dev/null +++ b/python/cibuildwheel_support/before_all_osx.sh @@ -0,0 +1,13 @@ +#! /bin/bash + +set -e +set -x + +brew install swig tree + +# Configure pip to use PyTorch extra-index-url for CPU +mkdir -p $HOME/.config/pip +echo "[global] +extra-index-url = https://download.pytorch.org/whl/cpu" > $HOME/.config/pip/pip.conf + +cp python/setup.py $HOME/setup.py.bkp \ No newline at end of file diff --git a/python/cibuildwheel_support/before_all_windows.sh b/python/cibuildwheel_support/before_all_windows.sh new file mode 100755 index 00000000..e8bb9266 --- /dev/null +++ b/python/cibuildwheel_support/before_all_windows.sh @@ -0,0 +1,110 @@ +#!/bin/bash + +set -e +set -x + +# Create pip directory if it doesn't exist +mkdir -p "C:\ProgramData\pip" + +# Create pip.ini file with PyTorch CPU index +echo "[global] +extra-index-url = https://download.pytorch.org/whl/cpu" > "C:\ProgramData\pip\pip.ini" + + +if [ "$ACCELERATOR" == "cu118" ]; then + curl --netrc-optional -L -nv -o cuda.exe https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_522.06_windows.exe + ./cuda.exe -s nvcc_11.8 nvrtc_11.8 nvrtc_dev_11.8 cudart_11.8 cufft_11.8 cufft_dev_11.8 cuda_profiler_api_11.8 nvtx_11.8 + rm cuda.exe + # Move CUDA folder to a path without spaces + mv "/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v11.8" /d/cuda + + # Create pip.ini file with PyTorch CUDA 11.8 index + echo "[global] +extra-index-url = https://download.pytorch.org/whl/cu118" > "C:\ProgramData\pip\pip.ini" + + pip install openmm-unofficial-cu11 torch==2.7.1 +elif [ "$ACCELERATOR" == "cu126" ]; then + curl --netrc-optional -L -nv -o cuda.exe https://developer.download.nvidia.com/compute/cuda/12.6.0/local_installers/cuda_12.6.0_560.76_windows.exe + ./cuda.exe -s nvcc_12.6 nvrtc_12.6 nvrtc_dev_12.6 cudart_12.6 cufft_12.6 cufft_dev_12.6 cuda_profiler_api_12.6 nvtx_12.6 + rm cuda.exe + # Move CUDA folder to a path without spaces + mv "/c/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.6" /d/cuda + # Create pip.ini file with PyTorch CUDA 12.6 index + echo "[global] +extra-index-url = https://download.pytorch.org/whl/cu126" > "C:\ProgramData\pip\pip.ini" + + pip install openmm-unofficial-cu12 torch==2.7.1 +elif [ "$ACCELERATOR" == "hip" ]; then + curl.exe --output HIP.exe --url https://download.amd.com/developer/eula/rocm-hub/AMD-Software-PRO-Edition-24.Q3-Win10-Win11-For-HIP.exe + ./HIP.exe -install + rm HIP.exe + pip install openmm-unofficial-cpu torch==2.7.1 +else + pip install openmm-unofficial-cpu torch==2.7.1 +fi + +# Download and extract OpenCL +curl --netrc-optional -L -nv -o OpenCL-SDK.zip https://github.com/KhronosGroup/OpenCL-SDK/releases/download/v2024.10.24/OpenCL-SDK-v2024.10.24-Win-x64.zip +unzip OpenCL-SDK.zip +OPENCL_PATH="$(pwd)/OpenCL-SDK-v2024.10.24-Win-x64" + + +###################################### +# Install dependencies with pip +PYTHONPREFIX=$(python -c 'import site; print(site.getsitepackages()[0])') +SITE_PACKAGES="$PYTHONPREFIX/Lib/site-packages/" + +CMAKE_FLAGS="-DENABLE_CUDA=OFF" +if [ "$ACCELERATOR" == "cu118" ] || [ "$ACCELERATOR" == "cu126" ]; then + ARCH_LIST=$(python -c "import torch; print(';'.join([f'{y[:-1]}.{y[-1]}' for y in [x[3:] for x in torch._C._cuda_getArchFlags().split() if x.startswith('sm_')]]))") + # CMakeLists.txt seems to ignore the CMAKE_CUDA_ARCHITECTURES variable, instead, it is overwritten by TORCH_CUDA_ARCH_LIST + ARCH_LIST_FMT=$(python -c "import torch; print(';'.join([y for y in [x[3:] for x in torch._C._cuda_getArchFlags().split() if x.startswith('sm_')]]))") + + export CUDA_PATH="/d/cuda" + CMAKE_FLAGS="-DTORCH_CUDA_ARCH_LIST=${ARCH_LIST}" + CMAKE_FLAGS+=" -DCMAKE_CUDA_ARCHITECTURES=${ARCH_LIST_FMT}" + CMAKE_FLAGS+=" -DCUDA_TOOLKIT_ROOT_DIR=${CUDA_PATH}" + CMAKE_FLAGS+=" -DCUDA_NVCC_EXECUTABLE=${CUDA_PATH}/bin/nvcc.exe" + CMAKE_FLAGS+=" -DCMAKE_CUDA_COMPILER=${CUDA_PATH}/bin/nvcc.exe" + CMAKE_FLAGS+=" -DCUDA_DRIVER_LIBRARY_PATH=${CUDA_PATH}/lib/x64/" +fi + +OPENCL_PATH="$(pwd)/OpenCL-SDK-v2024.10.24-Win-x64" + +# Configure build with Cmake +mkdir -p build +mkdir -p install +cd build + +echo $CMAKE_FLAGS +export LD_LIBRARY_PATH=/usr/lib/:/usr/local/cuda/targets/x86_64-linux/lib/:$LD_LIBRARY_PATH + +cmake .. -G "NMake Makefiles JOM" \ + -DCMAKE_INSTALL_PREFIX=../install \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_PREFIX_PATH=${SITE_PACKAGES} \ + -DUSE_SYSTEM_NVTX=1 \ + -DCMAKE_CXX_COMPILER=cl.exe \ + -DCMAKE_C_COMPILER=cl.exe \ + -DOPENMM_DIR=${SITE_PACKAGES}/openmm \ + -DPYTORCH_DIR=${SITE_PACKAGES}/torch \ + -DTorch_DIR=${SITE_PACKAGES}/torch/share/cmake/Torch \ + -DNN_BUILD_OPENCL_LIB=ON \ + -DOPENCL_INCLUDE_DIR="${OPENCL_PATH}/include" \ + -DOPENCL_LIBRARY="${OPENCL_PATH}/lib/OpenCL.lib" \ + $CMAKE_FLAGS + +# Build OpenMMTorch +jom -j 4 install +jom -j 4 PythonInstall + +cd .. + +cp build/python/setup.py python/ +cp build/python/openmmtorch.py python/openmmtorch/ +cp build/python/TorchPluginWrapper.cpp python/openmmtorch/ +cp -r install/include python/openmmtorch/ + +# Copy the libraries of openmm +mkdir -p python/openmm/ +cp -r install/lib/ python/openmm/ \ No newline at end of file diff --git a/python/cibuildwheel_support/before_build_osx.sh b/python/cibuildwheel_support/before_build_osx.sh new file mode 100644 index 00000000..60138099 --- /dev/null +++ b/python/cibuildwheel_support/before_build_osx.sh @@ -0,0 +1,54 @@ +#! /bin/bash + +set -e +set -x + +# Cleanup the python folder from previous build +cp $HOME/setup.py.bkp python/setup.py +rm -f python/TorchPluginWrapper.cpp +rm -f python/openmmtorch.py +rm -rf python/openmm/ +rm -rf python/openmmtorch/include + +# Install dependencies with pip +pip uninstall torch==2.7.1 openmm-unofficial-cpu -y +pip install torch==2.7.1 openmm-unofficial-cpu +unset SITE_PACKAGES +SITE_PACKAGES=$(python -c 'import site; print(site.getsitepackages()[0])') + +# Configure build with Cmake +rm -rf install build +mkdir -p build +mkdir -p install +cd build + +CC=/usr/bin/clang +CXX=/usr/bin/clang++ +CMAKE_CXX_FLAGS="-mmacosx-version-min=10.7" +CMAKE_SHARED_LINKER_FLAGS="-mmacosx-version-min=10.7" + +cmake .. \ + -DCMAKE_INSTALL_PREFIX=../install \ + -DCMAKE_BUILD_TYPE=Release \ + -DOPENMM_DIR=${SITE_PACKAGES}/openmm \ + -DPYTORCH_DIR=${SITE_PACKAGES}/torch \ + -DTorch_DIR=${SITE_PACKAGES}/torch/share/cmake/Torch \ + -DNN_BUILD_OPENCL_LIB=ON + +# Build OpenMMTorch +make -j4 install +make -j4 PythonInstall + +cd .. + +# Copy the generated python files and headers to openmmtorch +cp build/python/setup.py python/ +cp build/python/openmmtorch.py python/openmmtorch/ +cp build/python/TorchPluginWrapper.cpp python/openmmtorch/ +cp -r install/include python/openmmtorch/ + +# Copy the libraries of openmm +mkdir -p python/openmm/ +cp -r install/lib python/openmm/lib + + diff --git a/python/openmmtorch.i b/python/openmmtorch.i index a1529377..e5dcd94f 100644 --- a/python/openmmtorch.i +++ b/python/openmmtorch.i @@ -4,7 +4,7 @@ if sys.platform == 'win32': import os import torch import openmm - openmmtorch_library_path = openmm.version.openmm_library_path + openmmtorch_library_path = openmm.openmm_library_path _path = os.environ['PATH'] os.environ['PATH'] = r'%(lib)s;%(lib)s\plugins;%(path)s' % {'lib': openmmtorch_library_path, 'path': _path} diff --git a/python/openmmtorch/__init__.py b/python/openmmtorch/__init__.py new file mode 100644 index 00000000..874671dc --- /dev/null +++ b/python/openmmtorch/__init__.py @@ -0,0 +1 @@ +from .openmmtorch import * \ No newline at end of file diff --git a/python/pyproject.toml b/python/pyproject.toml new file mode 100644 index 00000000..f3d269a7 --- /dev/null +++ b/python/pyproject.toml @@ -0,0 +1,66 @@ +[project] +name = "PLACEHOLDER" +description = "Python wrapper for OpenMM (a C++ MD package)" +authors = [{ name = "Peter Eastman", email = "peastman@stanford.edu" }] +readme = "README.md" +license = "PSF-2.0" +requires-python = ">=3.8" +dynamic = ["version", "dependencies"] +classifiers = ["Programming Language :: Python :: 3"] + +[project.urls] +"Homepage" = "https://openmm.org" +"Bug Tracker" = "https://github.com/openmm/openmm/issues" +"Download URL" = "https://openmm.org" + +[tool.setuptools.packages.find] +include = ["openmmtorch*", "openmm*"] + +# [tool.setuptools.exclude-package-data] +# "*" = ["*.c", "*.cpp", "*.h", "*.cuh", "*.cu", ".gitignore"] + +[tool.pytest.ini_options] +python_files = "*.py" +python_classes = "Test" +python_functions = "test*" + +[build-system] +requires = ["setuptools>=78", "setuptools-scm>=8", "numpy", "cython"] +build-backend = "setuptools.build_meta" + + +[tool.cibuildwheel] +skip = [ + "cp38-*", + "pp*", + "*win32", + "*armv7l", + "*_i686", + "*_ppc64le", + "*_s390x", + "*_universal2", + "*-musllinux_*", +] +test-requires = ["pytest", "openmm-unofficial-cpu", "torch==2.7.1"] +test-command = "pytest {project}/python/tests --ignore={project}/python/tests/TestInteroperability.py" +manylinux-x86_64-image = "manylinux_2_28" +manylinux-aarch64-image = "manylinux_2_28" +environment-pass = ["ACCELERATOR"] +# container-engine = "docker; create_args: --gpus all" + +[tool.cibuildwheel.windows] +before-all = "bash {project}/python/cibuildwheel_support/before_all_windows.sh" +test-command = "pytest {project}\\python\\tests --ignore={project}\\python\\tests\\TestInteroperability.py" + +[tool.cibuildwheel.linux] +before-all = "bash {project}/python/cibuildwheel_support/before_all_linux.sh" +repair-wheel-command = [ + "auditwheel repair --exclude libcudart.so.* --exclude libOpenMMTorch.so --exclude libc10.so --exclude libc10_cuda.so --exclude libtorch.so --exclude libtorch_cuda.so --exclude libtorch_cpu.so --exclude libtorch_python.so --exclude libOpenMM.so --exclude libOpenMMCUDA.so --exclude libOpenMMHIP.so --exclude libOpenMMOpenCL.so --exclude libOpenMMDrude.so --exclude libOpenMMAmoeba.so --exclude libOpenMMRPMD.so --exclude libOpenCL.so.1 --exclude libcuda.so.1 --exclude libcufft.so.11 --exclude libcufft.so.10 --exclude libnvrtc.so.11.2 --exclude libnvrtc.so.12 --exclude libhiprtc.so.6 --exclude libamdhip64.so.6 -w {dest_dir} {wheel}", +] + +[tool.cibuildwheel.macos] +before-all = "bash {project}/python/cibuildwheel_support/before_all_osx.sh" +before-build = "bash {project}/python/cibuildwheel_support/before_build_osx.sh" +repair-wheel-command = [ + "delocate-wheel --ignore-missing-dependencies --require-archs {delocate_archs} -w {dest_dir} -v {wheel}", +] diff --git a/python/setup.py b/python/setup.py index 1a20eb43..4a2e55c9 100644 --- a/python/setup.py +++ b/python/setup.py @@ -9,28 +9,46 @@ nn_plugin_library_dir = '@NN_PLUGIN_LIBRARY_DIR@' torch_dir, _ = os.path.split('@TORCH_LIBRARY@') -extra_compile_args = ['-std=c++17'] +def _replace_name(name): + import pathlib + + pyproject_path = pathlib.Path(__file__).parent / "pyproject.toml" + with open(pyproject_path, 'r') as f: + pyproject_text = f.read() + pyproject_text = pyproject_text.replace("PLACEHOLDER", name) + with open(pyproject_path, 'w') as f: + f.write(pyproject_text) + +if os.getenv("ACCELERATOR", "").startswith("cpu"): + _replace_name("openmm-torch-unofficial-cpu") +if os.getenv("ACCELERATOR", "").startswith("cu"): + cuda_ver = os.getenv("ACCELERATOR", "")[2:4] + _replace_name(f"openmm-torch-unofficial-cu{cuda_ver}") + +extra_compile_args = ['-std=c++17', '-D_GLIBCXX_USE_CXX11_ABI=1'] extra_link_args = [] -libraries = ['OpenMM', 'OpenMMTorch'] -runtime_library_dirs = [os.path.join(openmm_dir, 'lib'), torch_dir] +libraries = ['OpenMM', 'OpenMMTorch', 'c10', 'torch', 'torch_cpu', 'torch_python'] +extra_deps = [] +if os.environ.get("ACCELERATOR", "").startswith("cu"): + libraries += ['c10_cuda', 'torch_cuda'] + cuda_ver = os.getenv("ACCELERATOR", "")[2:4] + extra_deps = [f'nvidia-cuda-runtime-cu{cuda_ver}'] + +runtime_library_dirs = ["$ORIGIN/../openmm/lib", "$ORIGIN/../torch/lib"] # For Windows change the compiler flag to /std:c++17 if platform.system() == 'Windows': extra_compile_args = ['/std:c++17'] - libraries += ['c10', 'torch'] - if os.environ.get("CUDA_HOME", None) is not None: - libraries += ['torch_cuda'] - else: - libraries += ['torch_cpu'] runtime_library_dirs = None # setup extra compile and link arguments on Mac if platform.system() == 'Darwin': extra_compile_args += ['-stdlib=libc++', '-mmacosx-version-min=10.13'] extra_link_args += ['-stdlib=libc++', '-mmacosx-version-min=10.13'] + runtime_library_dirs = ['@loader_path/../openmm/lib', '@loader_path/../torch/lib'] -extension = Extension(name='_openmmtorch', - sources=['TorchPluginWrapper.cpp'], +extension = Extension(name='openmmtorch._openmmtorch', + sources=['openmmtorch/TorchPluginWrapper.cpp'], libraries=libraries, include_dirs=[os.path.join(openmm_dir, 'include'), nn_plugin_header_dir] + torch_include_dirs, library_dirs=[os.path.join(openmm_dir, 'lib'), nn_plugin_library_dir, torch_dir], @@ -39,9 +57,9 @@ extra_link_args=extra_link_args ) -setup(name='openmmtorch', +setup( version=version, py_modules=['openmmtorch'], ext_modules=[extension], - install_requires=['openmm', 'torch'] + install_requires=extra_deps ) diff --git a/python/tests/TestInteroperability.py b/python/tests/TestInteroperability.py index 1c67c8a1..69236165 100644 --- a/python/tests/TestInteroperability.py +++ b/python/tests/TestInteroperability.py @@ -36,7 +36,7 @@ def forward(self, positions): system.addParticle(1.0) positions = pt.tensor([[-5, 0.0, 0.0], [5, 0.0, 0.0]], requires_grad=True) - with NamedTemporaryFile() as model_file: + with NamedTemporaryFile(delete=False) as model_file: # Save the model pt.jit.script(Model()).save(model_file.name) diff --git a/python/tests/TestTorchForce.py b/python/tests/TestTorchForce.py index 6eefed6e..d1a40a15 100644 --- a/python/tests/TestTorchForce.py +++ b/python/tests/TestTorchForce.py @@ -6,9 +6,32 @@ import torch as pt from tempfile import NamedTemporaryFile +class Central(pt.nn.Module): + def forward(self, pos): + return pos.pow(2).sum() + +class Forces(pt.nn.Module): + def forward(self, pos): + return pos.pow(2).sum(), -2 * pos + +class Global(pt.nn.Module): + def forward(self, pos, k): + return k * pos.pow(2).sum() + +class Periodic(pt.nn.Module): + def forward(self, pos, box): + box = box.diagonal().unsqueeze(0) + pos = pos - (pos / box).floor() * box + return pos.pow(2).sum() + +pt.jit.script(Central()).save('central.pt') +pt.jit.script(Forces()).save('forces.pt') +pt.jit.script(Global()).save('global.pt') +pt.jit.script(Periodic()).save('periodic.pt') + @pytest.mark.parametrize('model_file,', - ['../../tests/central.pt', - '../../tests/forces.pt']) + ['central.pt', + 'forces.pt']) def testConstructors(model_file): force = ot.TorchForce(model_file) model = pt.jit.load(model_file) @@ -17,9 +40,9 @@ def testConstructors(model_file): force = ot.TorchForce(model) @pytest.mark.parametrize('model_file, output_forces, use_module_constructor', - [('../../tests/central.pt', False, False,), - ('../../tests/forces.pt', True, False), - ('../../tests/forces.pt', True, True)]) + [('central.pt', False, False,), + ('forces.pt', True, False), + ('forces.pt', True, True)]) @pytest.mark.parametrize('use_cv_force', [True, False]) @pytest.mark.parametrize('platform', [mm.Platform.getPlatform(i).getName() for i in range(mm.Platform.getNumPlatforms())]) def testForce(model_file, output_forces, use_module_constructor, use_cv_force, platform): @@ -90,7 +113,7 @@ def forward(self, positions): assert pt.allclose(positions, self.positions) return pt.sum(positions) - with NamedTemporaryFile() as fd: + with NamedTemporaryFile(delete=False) as fd: numParticles = 10 system = mm.System() @@ -122,7 +145,7 @@ def forward(self, positions): def testProperties(): """ Test that the properties are correctly set and retrieved """ - force = ot.TorchForce('../../tests/central.pt') + force = ot.TorchForce('central.pt') force.setProperty('useCUDAGraphs', 'true') assert force.getProperties()['useCUDAGraphs'] == 'true' force.setProperty('useCUDAGraphs', 'false')