Add dummy cuda extension build (#521)

* Add dummy cuda extension build

* save

* run tests

* save

* fix

* pytorch cu110

* no vision

* save

* nl

* save

* save

* save

* save

* save

* save

* save

* autopep8

* more autopep8

* optional=True

* no todo

* fix

* refactor ci

* mypy

* save

* flake8

* os error

* Run cuaev test

* save

* clang format

* Delete delete-me.cu

* Search for cub in pytorch include paths

* nl

* instructions on test

* Add better handle cuaev install/not install

* test_cuaev.py

* clean

* fix
Co-authored-by: Farhad Ramezanghorbani <farhadrgh@users.noreply.github.com>

.github/workflows/cuda.yml

+name: cuda
+
+on:
+  pull_request:
+  push:
+    branches:
+      - master
+  schedule:
+    - cron: '0 0 * * *'
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      max-parallel: 4
+      matrix:
+        python-version: [3.8]
+
+    steps:
+    - uses: actions/checkout@v1
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v1
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Fetch submodules
+      run: git submodule update --init
+    - name: Install cuda
+      run: ci/install_cuda.sh
+    - name: Install dependencies
+      run: ci/install_dependencies.sh
+    - name: Install TorchANI
+      run: python setup.py install --cuaev
+    - name: Run cuaev test
+      run: pytest tests/test_cuaev.py -v

.github/workflows/deploy-docs.yml

@@ -20,10 +20,9 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
-      run: |
-        ci/install_dependencies.sh
-        pip install h5py tb-nightly sphinx sphinx_rtd_theme matplotlib pillow sphinx-gallery
-        pip install .
+      run: ci/install_dependencies.sh
+    - name: Install TorchANI
+      run: python setup.py install
     - name: Download data files
       run: ./download.sh
     # double check that this is indeed run under a release

.github/workflows/deploy-pypi.yml

@@ -22,9 +22,7 @@ jobs:
     - name: Fail build on non-release commits
       run: git describe --exact-match --tags HEAD
     - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install twine wheel
+      run: ci/install_dependencies.sh
     - name: Deploy
       run: |
         rm -rf dist/*

.github/workflows/deploy-test-pypi.yml

@@ -24,9 +24,7 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install twine wheel
+      run: ci/install_dependencies.sh
     - name: Deploy
       run: |
         rm -rf dist/*

.github/workflows/docs.yml

@@ -24,10 +24,9 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
-      run: |
-        ci/install_dependencies.sh
-        pip install h5py tb-nightly sphinx sphinx_rtd_theme matplotlib pillow sphinx-gallery
-        pip install .
+      run: ci/install_dependencies.sh
+    - name: Install TorchANI
+      run: python setup.py install
     - name: Download data files
       run: ./download.sh
     # build docs

.github/workflows/flake8.yml

@@ -25,9 +25,7 @@ jobs:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
       run: |
-        python -m pip install --upgrade pip
-        pip install .
-    - name: Lint with flake8
-      run: |
+        pip install --upgrade pip
         pip install flake8
-        flake8 .
+    - name: Lint with flake8
+      run: flake8 .

.github/workflows/mypy.yml

@@ -23,8 +23,9 @@ jobs:
       uses: actions/setup-python@v1
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Type checking with mypy
+    - name: Install dependencies
       run: |
         pip install --upgrade pip
         pip install mypy
-        mypy --no-site-packages --ignore-missing-imports torchani/{aev,nn,utils,models,ase}.py
+    - name: Type checking with mypy
+      run: mypy --no-site-packages --ignore-missing-imports torchani/{aev,nn,utils,models,ase}.py

.github/workflows/runnable_submodules.yml

@@ -24,10 +24,9 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install dependencies
-      run: |
-        ci/install_dependencies.sh
-        pip install h5py .
-        pip install .
+      run: ci/install_dependencies.sh
+    - name: Install TorchANI
+      run: python setup.py install
     - name: Download data files
       run: ./download.sh
     - name: Run submodules

.github/workflows/stable-torch.yml

 name: stable-torch
 
 on:
+  pull_request:
   push:
     branches:
       - master
@@ -25,11 +26,9 @@ jobs:
     - name: Fetch submodules
       run: git submodule update --init
     - name: Install dependencies
-      run: |
-        pip install --upgrade pip
-        pip install torch pytest
-        pip install -r test_requirements.txt
-        pip install .
+      run: ci/install_dependencies.sh
+    - name: Install TorchANI
+      run: python setup.py install
     - name: Download data files
       run: ./download.sh
     - name: Run tests

.github/workflows/tools.yml

@@ -26,15 +26,13 @@ jobs:
     - name: Fetch submodules
       run: git submodule update --init
     - name: Install dependencies
-      run: |
-        ci/install_dependencies.sh
-        pip install .
+      run: ci/install_dependencies.sh
+    - name: Install TorchANI
+      run: python setup.py install
     - name: Download data files
       run: ./download.sh
     - name: Inference benchmark
       run: python tools/inference-benchmark.py --tqdm dataset/xyz_files/CH4-5.xyz
-    - name: Install more dependencies
-      run: pip install h5py
     - name: COMP6 Benchmark
       run: python tools/comp6.py dataset/COMP6/COMP6v1/s66x8
     - name: Training Benchmark

.github/workflows/unittest.yml

@@ -21,7 +21,7 @@ jobs:
           test_utils.py, test_ase.py, test_energies.py, test_periodic_table_indexing.py,
           test_neurochem.py, test_vibrational.py, test_ensemble.py, test_padding.py,
           test_data.py, test_forces.py, test_structure_optim.py, test_jit_builtin_models.py, 
-          test_grad.py]
+          test_grad.py, test_cuaev.py]
 
     steps:
     - uses: actions/checkout@v1
@@ -32,10 +32,9 @@ jobs:
     - name: Fetch submodules
       run: git submodule update --init
     - name: Install dependencies
-      run: |
-        ci/install_dependencies.sh
-        pip install -r test_requirements.txt
-        pip install .
+      run: ci/install_dependencies.sh
+    - name: Install TorchANI
+      run: python setup.py install
     - name: Download data files
       run: ./download.sh
     - name: Run single test

.gitignore

@@ -38,3 +38,4 @@ Untitled.ipynb
 /test*.py
 .coverage
 htmlcov/
+/include

README.md

@@ -89,10 +89,14 @@ pip install -e .
 After TorchANI has been installed, you can build the documents by running `sphinx-build docs build`. But make sure you
 install dependencies:
 ```bash
-pip install sphinx sphinx-gallery pillow matplotlib sphinx_rtd_theme
+pip install -r docs_requirements.txt
 ```
 
-To manually run unit tests, do `python setup.py nosetests`
+To manually run unit tests, do
+
+```bash
+pytest -v
+```
 
 If you opened a pull request, you could see your generated documents at https://aiqm.github.io/torchani-test-docs/ after you `docs` check succeed.
 Keep in mind that this repository is only for the purpose of convenience of development, and only keeps the latest push.

ci/install_cuda.sh

+#!/bin/bash
+
+# command copy-pasted from https://developer.nvidia.com/cuda-downloads?target_os=Linux&target_arch=x86_64&target_distro=Ubuntu&target_version=1804&target_type=debnetwork
+
+wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-ubuntu1804.pin
+sudo mv cuda-ubuntu1804.pin /etc/apt/preferences.d/cuda-repository-pin-600
+sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/7fa2af80.pub
+sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/ /"
+sudo apt-get update
+sudo apt-get -y install cuda

ci/install_dependencies.sh

 #!/bin/bash
 
 pip install --upgrade pip
-pip install --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
+pip install twine wheel
+pip install --pre torch -f https://download.pytorch.org/whl/nightly/cu110/torch_nightly.html
 pip install -r test_requirements.txt
+pip install -r docs_requirements.txt

docs_requirements.txt

+h5py
+tb-nightly
+sphinx
+sphinx_rtd_theme
+matplotlib
+pillow
+sphinx-gallery

setup.py

+import os
+import subprocess
 from setuptools import setup, find_packages
+from distutils import log
+import sys
+
+BUILD_CUAEV = '--cuaev' in sys.argv
+if BUILD_CUAEV:
+    sys.argv.remove('--cuaev')
+
+if not BUILD_CUAEV:
+    log.warn("Will not install cuaev")
 
 with open("README.md", "r") as fh:
     long_description = fh.read()
 
-setup_attrs = {
-    'name': 'torchani',
-    'description': 'PyTorch implementation of ANI',
-    'long_description': long_description,
-    'long_description_content_type': "text/markdown",
-    'url': 'https://github.com/aiqm/torchani',
-    'author': 'Xiang Gao',
-    'author_email': 'qasdfgtyuiop@gmail.com',
-    'license': 'MIT',
-    'packages': find_packages(),
-    'include_package_data': True,
-    'use_scm_version': True,
-    'setup_requires': ['setuptools_scm'],
-    'install_requires': [
+
+def maybe_download_cub():
+    import torch
+    dirs = torch.utils.cpp_extension.include_paths(cuda=True)
+    for d in dirs:
+        cubdir = os.path.join(d, 'cub')
+        log.info(f'Searching for cub at {cubdir}...')
+        if os.path.isdir(cubdir):
+            log.info(f'Found cub in {cubdir}')
+            return []
+    # if no cub, download it to include dir from github
+    if not os.path.isdir('./include/cub'):
+        if not os.path.exists('./include'):
+            os.makedirs('include')
+        commands = """
+        echo "Downloading CUB library";
+        wget -q https://github.com/NVIDIA/cub/archive/main.zip;
+        unzip -q main.zip -d include;
+        mv include/cub-main/cub include;
+        echo "Removing unnecessary files";
+        rm main.zip;
+        rm -rf include/cub-main;
+        """
+        subprocess.run(commands, shell=True, check=True, universal_newlines=True)
+    return [os.path.abspath("./include")]
+
+
+def cuda_extension():
+    import torch
+    from torch.utils.cpp_extension import CUDAExtension
+
+    nvcc_args = ["-gencode=arch=compute_50,code=sm_50", "-gencode=arch=compute_60,code=sm_60",
+                 "-gencode=arch=compute_61,code=sm_61", "-gencode=arch=compute_70,code=sm_70",
+                 "-Xptxas=-v", '--expt-extended-lambda', '-use_fast_math']
+    cuda_version = float(torch.version.cuda)
+    if cuda_version >= 10:
+        nvcc_args.append("-gencode=arch=compute_75,code=sm_75")
+    if cuda_version >= 11:
+        nvcc_args.append("-gencode=arch=compute_80,code=sm_80")
+    if cuda_version >= 11.1:
+        nvcc_args.append("-gencode=arch=compute_86,code=sm_86")
+    return CUDAExtension(
+        name='_real_cuaev',
+        pkg='torchani.cuaev._real_cuaev',
+        sources=['torchani/cuaev/aev.cu'],
+        include_dirs=maybe_download_cub(),
+        extra_compile_args={'nvcc': nvcc_args},
+        optional=True)
+
+
+def cuaev_kwargs():
+    if not BUILD_CUAEV:
+        return dict(
+            provides=['torchani']
+        )
+    from torch.utils.cpp_extension import BuildExtension
+    kwargs = dict(
+        provides=[
+            'torchani',
+            'torchani.cuaev',
+        ],
+        ext_modules=[
+            cuda_extension()
+        ],
+        cmdclass={
+            'build_ext': BuildExtension,
+        })
+    return kwargs
+
+
+setup(
+    name='torchani',
+    description='PyTorch implementation of ANI',
+    long_description=long_description,
+    long_description_content_type="text/markdown",
+    url='https://github.com/aiqm/torchani',
+    author='Xiang Gao',
+    author_email='qasdfgtyuiop@gmail.com',
+    license='MIT',
+    packages=find_packages(),
+    include_package_data=True,
+    use_scm_version=True,
+    setup_requires=['setuptools_scm'],
+    install_requires=[
         'torch',
         'lark-parser',
-        'requests'
+        'requests',
+        'importlib_metadata',
     ],
-}
-
-setup(**setup_attrs)
+    **cuaev_kwargs()
+)

test_requirements.txt

@@ -5,3 +5,4 @@ h5py
 pillow
 pkbar
 pyyaml
+pytest

tests/test_cuaev.py

+import torchani
+import unittest
+import torch
+
+skipIfNoGPU = unittest.skipIf(not torch.cuda.is_available(),
+                              'There is no device to run this test')
+
+
+@unittest.skipIf(torchani.cuaev.is_installed, "only valid when cuaev not installed")
+class TestCUAEVNotInstalled(unittest.TestCase):
+
+    def testCuComputeAEV(self):
+        self.assertRaisesRegex(RuntimeError, "cuaev is not installed", lambda: torchani.cuaev.cuComputeAEV())
+
+
+@unittest.skipIf(not torchani.cuaev.is_installed, "only valid when cuaev is installed")
+class TestCUAEV(unittest.TestCase):
+
+    @skipIfNoGPU
+    def testHello(self):
+        # TODO: this should be removed when a real cuaev is merged
+        self.assertEqual("Hello World!!!", torchani.cuaev.cuComputeAEV())
+
+
+if __name__ == '__main__':
+    unittest.main()

torchani/__init__.py

@@ -36,6 +36,7 @@ from . import utils
 from . import neurochem
 from . import models
 from . import units
+from . import cuaev
 from pkg_resources import get_distribution, DistributionNotFound
 
 try:
@@ -45,7 +46,7 @@ except DistributionNotFound:
     pass
 
 __all__ = ['AEVComputer', 'EnergyShifter', 'ANIModel', 'Ensemble', 'SpeciesConverter',
-           'utils', 'neurochem', 'models', 'units']
+           'utils', 'neurochem', 'models', 'units', 'cuaev']
 
 try:
     from . import ase  # noqa: F401