diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 760f69b..cfe9748 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -114,6 +114,22 @@ jobs: conda activate nnpops conda list + - name: Pip build (without CUDA) + if: ${{ !matrix.enable_cuda }} + run: | + conda activate nnpops + sed -i -e '/"torch/c "torch==${{ matrix.pytorch }}",' pyproject.toml + pip -vvv wheel --wheel-dir=$(pwd) . --extra-index-url https://download.pytorch.org/whl/cpu + rm -rf build + + - name: Pip build (with CUDA) + if: ${{ matrix.enable_cuda }} + run: | + conda activate nnpops + sed -i -e '/"torch/c "torch==${{ matrix.pytorch }}",' pyproject.toml + pip -vvv wheel --wheel-dir=$(pwd) . + rm -rf build + - name: Configure, compile, and install run: | conda activate nnpops @@ -129,3 +145,4 @@ jobs: conda activate nnpops cd build ctest --verbose --exclude-regex TestCuda + diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml new file mode 100644 index 0000000..14d0ef5 --- /dev/null +++ b/.github/workflows/wheels.yml @@ -0,0 +1,43 @@ +name: Build + +on: [push, pull_request] + +jobs: + build_wheels: + name: Build wheels on cp${{ matrix.python }}-linux + + runs-on: ${{ matrix.os }} + strategy: + matrix: + os: [ubuntu-latest] #, windows-2019, macos-11] + python: ['3.10'] + torch: ['1.12'] + cuda: ['11'] + steps: + - uses: actions/checkout@v3 + - name: Setup Python + uses: actions/setup-python@v4 + with: + python-version: ${{ matrix.python }} + - name: Build wheels + uses: pypa/cibuildwheel@v2.13.1 + env: + CIBW_SKIP: "cp36-* *-win32 *i686" + + CIBW_ENVIRONMENT: TORCH_VERSION=${{ matrix.torch }} CUDA_VERSION=${{ matrix.cuda }} + with: + package-dir: . + output-dir: wheelhouse + + + # env: + # CIBW_SOME_OPTION: value + # ... + # with: + # package-dir: . + # output-dir: wheelhouse + # config-file: "{package}/pyproject.toml" + + # - uses: actions/upload-artifact@v3 + # with: + # path: ./wheelhouse/*.whl diff --git a/CMakeLists.txt b/CMakeLists.txt index 5529ff6..b3df303 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -9,7 +9,7 @@ if(ENABLE_CUDA) endif(ENABLE_CUDA) # Find dependencies -find_package(Python3 REQUIRED COMPONENTS Interpreter Development) +find_package(Python3 REQUIRED COMPONENTS Interpreter Development.Module) find_package(Torch REQUIRED) enable_testing() diff --git a/environment.yml b/environment.yml index 9db9ae2..26b5612 100644 --- a/environment.yml +++ b/environment.yml @@ -12,3 +12,4 @@ dependencies: - python 3.10.* - pytorch-gpu 2.0.* - sysroot_linux-64 2.17 + - pip diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..61920ed --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,27 @@ +[build-system] +requires = [ + "setuptools>=42", + "wheel", + "ninja", + "cmake>=3.23", + "GitPython", + "nvidia-pyindex", + "torch>=1.11", + "Cython" + +] +build-backend = "setuptools.build_meta" +python_requires = ">=3.7" + +[project] +name = "nnpops" +version = "0.5" +dependencies = [ + "numpy", + "torch>=1.11", + "torchani>=2.2", + "mdtraj" +] + +[project.optional-dependencies] +test = ["pytest>=5.0"] \ No newline at end of file diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..1a1cb81 --- /dev/null +++ b/setup.py @@ -0,0 +1,178 @@ +import os +import re +import subprocess +import sys +from pathlib import Path +#import git +from setuptools import Extension, setup +from setuptools.command.build_ext import build_ext +import torch + +from typing import Optional, Dict +# Convert distutils Windows platform specifiers to CMake -A arguments +PLAT_TO_CMAKE = { + "win32": "Win32", + "win-amd64": "x64", + "win-arm32": "ARM", + "win-arm64": "ARM64", +} + +# A CMakeExtension needs a sourcedir instead of a file list. +# The name must be the _single_ output extension from the CMake build. +# If you need multiple extensions, see scikit-build. +class CMakeExtension(Extension): + def __init__(self, name: str, sourcedir: str = "", extra_args: Optional[Dict[str, str]] = None) -> None: + super().__init__(name, sources=[]) + self.sourcedir = os.fspath(Path(sourcedir).resolve()) + #Store a list of extra arguments to pass to CMake, prepend -D to each + if extra_args is not None: + print("Extra args: ", extra_args) + self.extra_args = [f"-D{key}={value}" for key, value in extra_args.items()] if extra_args else [] + + + +class CMakeBuild(build_ext): + def build_extension(self, ext: CMakeExtension) -> None: + # Must be in this form due to bug in .resolve() only fixed in Python 3.10+ + ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name) + extdir = ext_fullpath.parent.resolve() + + # Using this requires trailing slash for auto-detection & inclusion of + # auxiliary "native" libs + + debug = int(os.environ.get("DEBUG", 0)) if self.debug is None else self.debug + cfg = "Debug" if debug else "Release" + + # CMake lets you override the generator - we need to check this. + # Can be set with Conda-Build, for example. + cmake_generator = os.environ.get("CMAKE_GENERATOR", "make") + + # Set Python_EXECUTABLE instead if you use PYBIND11_FINDPYTHON + # EXAMPLE_VERSION_INFO shows you how to pass a value into the C++ code + # from Python. + cmake_args = [ + f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}{os.sep}{ext.name}", + f"-DPYTHON_EXECUTABLE={sys.executable}", + f"-DCMAKE_BUILD_TYPE={cfg}", # not used on MSVC, but no harm + ] + build_args = [] + # Adding CMake arguments set as environment variable + # (needed e.g. to build for ARM OSx on conda-forge) + if "CMAKE_ARGS" in os.environ: + cmake_args += [item for item in os.environ["CMAKE_ARGS"].split(" ") if item] + + # In this example, we pass in the version to C++. You might not need to. + cmake_args += [f"-DEXAMPLE_VERSION_INFO={self.distribution.get_version()}"] + + if self.compiler.compiler_type != "msvc": + # Using Ninja-build since it a) is available as a wheel and b) + # multithreads automatically. MSVC would require all variables be + # exported for Ninja to pick it up, which is a little tricky to do. + # Users can override the generator with CMAKE_GENERATOR in CMake + # 3.15+. + if not cmake_generator or cmake_generator == "Ninja": + try: + import ninja + + ninja_executable_path = Path(ninja.BIN_DIR) / "ninja" + cmake_args += [ + "-GNinja", + f"-DCMAKE_MAKE_PROGRAM:FILEPATH={ninja_executable_path}", + ] + except ImportError: + pass + + else: + # Single config generators are handled "normally" + single_config = any(x in cmake_generator for x in {"NMake", "Ninja"}) + + # CMake allows an arch-in-generator style for backward compatibility + contains_arch = any(x in cmake_generator for x in {"ARM", "Win64"}) + + # Specify the arch if using MSVC generator, but only if it doesn't + # contain a backward-compatibility arch spec already in the + # generator name. + if not single_config and not contains_arch: + cmake_args += ["-A", PLAT_TO_CMAKE[self.plat_name]] + + # Multi-config generators have a different way to specify configs + if not single_config: + cmake_args += [ + f"-DCMAKE_LIBRARY_OUTPUT_DIRECTORY_{cfg.upper()}={extdir}" + ] + build_args += ["--config", cfg] + + if sys.platform.startswith("darwin"): + # Cross-compile support for macOS - respect ARCHFLAGS if set + archs = re.findall(r"-arch (\S+)", os.environ.get("ARCHFLAGS", "")) + if archs: + cmake_args += ["-DCMAKE_OSX_ARCHITECTURES={}".format(";".join(archs))] + cmake_args += ext.extra_args + # Set CMAKE_BUILD_PARALLEL_LEVEL to control the parallel build level + # across all generators. + if "CMAKE_BUILD_PARALLEL_LEVEL" not in os.environ: + # self.parallel is a Python 3 only way to set parallel jobs by hand + # using -j in the build_ext call, not supported by pip or PyPA-build. + if hasattr(self, "parallel") and self.parallel: + # CMake 3.12+ only. + build_args += [f"-j{self.parallel}"] + + build_temp = Path(self.build_temp) / ext.name + if not build_temp.exists(): + build_temp.mkdir(parents=True) + + subprocess.run( + ["cmake", ext.sourcedir, *cmake_args], cwd=build_temp, check=True + ) + subprocess.run( + ["cmake", "--build", ".", *build_args], cwd=build_temp, check=True + ) + +extra_args = {} + +if "CC" in os.environ: + extra_args["CMAKE_C_COMPILER"] = os.environ.get("CC", "") +if "CXX" in os.environ: + extra_args["CMAKE_CXX_COMPILER"] = os.environ.get("CXX", "") +# if torch.backends.cuda.is_built(): +# extra_args["ENABLE_CUDA"] = "ON" +# ARCHES = [52, 60, 61, 70] +# DEPRECATED_IN_11 = [35, 50] +# cuda_version_major= int(torch.version.cuda.split(".")[0]) +# cuda_version_minor= int(torch.version.cuda.split(".")[1]) +# if cuda_version_major >= 11 or (cuda_version_major == 11 and cuda_version_minor >= 1): +# LATEST_ARCH = 90 +# ARCHES += [75, 80, 86] +# elif cuda_version_major == 11 and cuda_version_minor >= 1: +# LATEST_ARCH = 86 +# ARCHES += [75, 80] +# elif cuda_version_major == 11 and cuda_version_minor >= 0: +# LATEST_ARCH = 80 +# ARCHES += [75] +# elif cuda_version_major >= 10: +# LATEST_ARCH = 75 +# ARCHES += DEPRECATED_IN_11 +# else: +# raise RuntimeError("Unsupported CUDA version") +# CMAKE_CUDA_ARCHS = ";".join([str(arch) for arch in ARCHES] + [f"{LATEST_ARCH}-real", f"{LATEST_ARCH}-virtual"]) +# extra_args["CMAKE_CUDA_ARCHITECTURES"] = "OFF" #CMAKE_CUDA_ARCHS +# else: +extra_args["CMAKE_CUDA_ARCHITECTURES"] = "OFF" +extra_args["ENABLE_CUDA"] = "OFF" + +extra_args["CMAKE_PREFIX_PATH"] = torch.utils.cmake_prefix_path +torch_version = os.environ.get("TORCH_VERSION", torch.__version__) +cuda_version = os.environ.get("CUDA_VERSION", torch.version.cuda) +#tag = git.Repo(search_parent_directories=True).git.describe("--tags", always=True) +#version = tag.lstrip('v').split('-')[0] +setup( + ext_modules=[CMakeExtension(name="NNPOps", sourcedir=".", extra_args=extra_args)], + cmdclass={"build_ext": CMakeBuild}, + packages=["NNPOps", "NNPOps.neighbors", "NNPOps.pme"], + package_dir={ + 'NNPOps': 'src/pytorch', + 'NNPOps.neighbors': 'src/pytorch/neighbors', + 'NNPOps.pme': 'src/pytorch/pme'}, + package_data={'NNPOps': ['lib/*.so', 'lib/*.dll', 'lib/*.dylib']}, + install_requires=[f"torch=={torch_version}", f"nvidia-cuda-nvcc-cu{cuda_version}"], +)