Python CMakeExtension

.github/workflows/install.yml

@@ -27,22 +27,11 @@ jobs:
       - name: Install PyTorch ${{ matrix.torch-version }}+cpu
         run: |
           pip install torch==${{ matrix.torch-version}}+cpu -f https://download.pytorch.org/whl/torch_stable.html
-      - name: Configure
-        run: |
-          Torch_DIR=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'` cmake -S . -B build -DCMAKE_BUILD_TYPE=Debug
-      - name: Build
-        run: |
-          cmake --build build --config Debug
-      - name: Test
-        run: |
-          cd build
-          ctest -C Debug --output-on-failure --verbose
-          cd ..
 
       - name: Install main package
         run: |
-          pip install .
+          pip install -e .
 
       - name: Test imports
         run: |
-          python -c "import pyg_lib"
+          python -c 'import pyg_lib;print(pyg_lib.cuda_version())'

.github/workflows/testing.yml

@@ -0,0 +1,43 @@
+name: Testing
+
+on:  # yamllint disable-line rule:truthy
+  push:
+    branches:
+      - master
+  pull_request:
+
+jobs:
+
+  gtest:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: [3.9]
+        torch-version: [1.11.0]
+
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install PyTorch ${{ matrix.torch-version }}+cpu
+        run: |
+          pip install torch==${{ matrix.torch-version}}+cpu -f https://download.pytorch.org/whl/torch_stable.html
+
+      - name: Configure
+        run: |
+          Torch_DIR=`python -c 'import torch;print(torch.utils.cmake_prefix_path)'` cmake -S . -B build -DCMAKE_BUILD_TYPE=DEBUG
+
+      - name: Build
+        run: |
+          cmake --build build --config DEBUG
+
+      - name: Run tests
+        run: |
+          cd build
+          ctest -C DEBUG --output-on-failure --verbose
+          cd ..

.gitignore

@@ -1,12 +1,13 @@
 __pycache__/
 .pytest_cache/
 build/
+Testing/
 dist/
 alpha/
 .cache/
 .eggs/
 *.egg-info/
-.ipynb_checkpoints
+*.so
 .coverage
 .coverage.*
 .vscode

CMakeLists.txt

@@ -3,8 +3,8 @@ project(pyg)
 set(CMAKE_CXX_STANDARD 14)
 set(PYG_VERSION 0.0.0)
 
+option(USE_PYTHON "Link to Python when building" OFF)
 option(WITH_CUDA "Enable CUDA support" OFF)
-option(USE_PYTHON "Link to Python when building" ON)
 option(BUILD_TEST "Enable testing" ON)
 
 if(WITH_CUDA)

pyg_lib/__init__.py

@@ -1,5 +1,20 @@
+import os.path as osp
+import importlib
+import torch
+
 __version__ = '0.0.0'
 
+loader_details = (importlib.machinery.ExtensionFileLoader,
+                  importlib.machinery.EXTENSION_SUFFIXES)
+
+path = osp.abspath(osp.join(osp.dirname(__file__), '..'))
+ext_finder = importlib.machinery.FileFinder(path, loader_details)
+spec = ext_finder.find_spec('libpyg')
+torch.ops.load_library(spec.origin)
+
+cuda_version = torch.ops.pyg.cuda_version
+
 __all__ = [
     '__version__',
+    'cuda_version',
 ]

setup.py

@@ -1,8 +1,59 @@
-from setuptools import find_packages, setup
+import importlib
+import os
+import os.path as osp
+import subprocess
+
+from setuptools import Extension, find_packages, setup
+from setuptools.command.build_ext import build_ext
 
 __version__ = '0.0.0'
 URL = 'https://github.com/pyg-team/pyg-lib'
 
+
+class CMakeExtension(Extension):
+    def __init__(self, name, sourcedir=''):
+        Extension.__init__(self, name, sources=[])
+        self.sourcedir = os.path.abspath(sourcedir)
+
+
+class CMakeBuild(build_ext):
+    def get_ext_filename(self, ext_name):
+        # Remove Python ABI suffix:
+        ext_filename = super().get_ext_filename(ext_name)
+        ext_filename_parts = ext_filename.split('.')
+        ext_filename_parts = ext_filename_parts[:-2] + ext_filename_parts[-1:]
+        return '.'.join(ext_filename_parts)
+
+    def build_extension(self, ext):
+        import torch
+
+        extdir = os.path.abspath(osp.dirname(self.get_ext_fullpath(ext.name)))
+
+        if self.debug is None:
+            self.debug = bool(int(os.environ.get('DEBUG', 0)))
+
+        if not osp.exists(self.build_temp):
+            os.makedirs(self.build_temp)
+
+        cmake_args = [
+            '-DUSE_PYTHON=ON',
+            f'-DCMAKE_LIBRARY_OUTPUT_DIRECTORY={extdir}',
+            f'-DCMAKE_BUILD_TYPE={"DEBUG" if self.debug else "RELEASE"}',
+            f'-DCMAKE_PREFIX_PATH={torch.utils.cmake_prefix_path}',
+            f'-DWITH_CUDA={"ON" if torch.cuda.is_available() else "OFF"}',
+        ]
+
+        if importlib.util.find_spec('ninja') is not None:
+            cmake_args += ['-GNinja']
+
+        build_args = []
+
+        subprocess.check_call(['cmake', ext.sourcedir] + cmake_args,
+                              cwd=self.build_temp)
+        subprocess.check_call(['cmake', '--build', '.'] + build_args,
+                              cwd=self.build_temp)
+
+
 install_requires = []
 
 test_requires = [
@@ -18,8 +69,8 @@
     name='pyg_lib',
     version=__version__,
     description='Low-Level Graph Neural Network Operators for PyG',
-    author='Matthias Fey',
-    author_email='[email protected]',
+    author='PyG Team',
+    author_email='[email protected]',
     url=URL,
     download_url=f'{URL}/archive/{__version__}.tar.gz',
     keywords=[
@@ -36,5 +87,7 @@
         'dev': dev_requires,
     },
     packages=find_packages(),
+    ext_modules=[CMakeExtension('libpyg')],
+    cmdclass={'build_ext': CMakeBuild},
     include_package_data=True,
 )