Merge pull request #40 from rusty1s/metis

[WIP] Partition

Author: rusty1s

.travis.yml

@@ -60,6 +60,7 @@ jobs:
 install:
   - source script/cuda.sh
   - source script/conda.sh
+  - source script/metis.sh
   - conda create --yes -n test python="${PYTHON_VERSION}"
   - source activate test
   - conda install pytorch=${TORCH_VERSION} ${TOOLKIT} -c pytorch --yes

README.md

@@ -59,17 +59,20 @@ $ echo $CPATH
 >>> /usr/local/cuda/include:...
 ```
 
+If you want to additionally build `torch-sparse` with METIS support, *e.g.* for partioning, please download and install the [METIS library](http://glaros.dtc.umn.edu/gkhome/metis/metis/download) by following the instructions in the `Install.txt` file.
+Afterwards, set the environment variable `WITH_METIS=1`.
+
 Then run:
 
 ```
 pip install torch-scatter torch-sparse
 ```
 
-When running in a docker container without nvidia driver, PyTorch needs to evaluate the compute capabilities and may fail.
+When running in a docker container without NVIDIA driver, PyTorch needs to evaluate the compute capabilities and may fail.
 In this case, ensure that the compute capabilities are set via `TORCH_CUDA_ARCH_LIST`, *e.g.*:
 
 ```
-export TORCH_CUDA_ARCH_LIST = "6.0 6.1 7.2+PTX 7.5+PTX"
+export TORCH_CUDA_ARCH_LIST="6.0 6.1 7.2+PTX 7.5+PTX"
 ```
 
 ## Functions

csrc/cpu/metis_cpu.cpp

@@ -0,0 +1,36 @@
+#include "metis_cpu.h"
+
+#ifdef WITH_METIS
+#include <metis.h>
+#endif
+
+#include "utils.h"
+
+torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
+                            int64_t num_parts, bool recursive) {
+#ifdef WITH_METIS
+  CHECK_CPU(rowptr);
+  CHECK_CPU(col);
+
+  int64_t nvtxs = rowptr.numel() - 1;
+  auto part = torch::empty(nvtxs, rowptr.options());
+
+  auto *xadj = rowptr.data_ptr<int64_t>();
+  auto *adjncy = col.data_ptr<int64_t>();
+  int64_t ncon = 1;
+  int64_t objval = -1;
+  auto part_data = part.data_ptr<int64_t>();
+
+  if (recursive) {
+    METIS_PartGraphRecursive(&nvtxs, &ncon, xadj, adjncy, NULL, NULL, NULL,
+                             &num_parts, NULL, NULL, NULL, &objval, part_data);
+  } else {
+    METIS_PartGraphKway(&nvtxs, &ncon, xadj, adjncy, NULL, NULL, NULL,
+                        &num_parts, NULL, NULL, NULL, &objval, part_data);
+  }
+
+  return part;
+#else
+  AT_ERROR("Not compiled with METIS support");
+#endif
+}

csrc/cpu/metis_cpu.h

@@ -0,0 +1,6 @@
+#pragma once
+
+#include <torch/extension.h>
+
+torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
+                            int64_t num_parts, bool recursive);

csrc/metis.cpp

@@ -0,0 +1,24 @@
+#include <Python.h>
+#include <torch/script.h>
+
+#include "cpu/metis_cpu.h"
+
+#ifdef _WIN32
+PyMODINIT_FUNC PyInit__metis(void) { return NULL; }
+#endif
+
+torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
+                        int64_t num_parts, bool recursive) {
+  if (rowptr.device().is_cuda()) {
+#ifdef WITH_CUDA
+    AT_ERROR("No CUDA version supported");
+#else
+    AT_ERROR("Not compiled with CUDA support");
+#endif
+  } else {
+    return partition_cpu(rowptr, col, num_parts, recursive);
+  }
+}
+
+static auto registry =
+    torch::RegisterOperators().op("torch_sparse::partition", &partition);

script/metis.sh

@@ -0,0 +1,29 @@
+#!/bin/bash
+
+METIS=metis-5.1.0
+export WITH_METIS=1
+
+wget -nv http://glaros.dtc.umn.edu/gkhome/fetch/sw/metis/${METIS}.tar.gz
+tar -xvzf ${METIS}.tar.gz
+cd ${METIS} || exit
+sed -i.bak -e 's/IDXTYPEWIDTH 32/IDXTYPEWIDTH 64/g' include/metis.h
+
+if [ "${TRAVIS_OS_NAME}" != "windows" ]; then
+  make config
+  make
+  sudo make install
+else
+  # Fix GKlib on Windows: https://github.com/jlblancoc/suitesparse-metis-for-windows/issues/6
+  sed -i.bak -e '61,69d' GKlib/gk_arch.h
+
+  cd build || exit
+
+  cmake .. -A x64  # Ensure we are building with x64
+  cmake --build . --config "Release" --target ALL_BUILD
+  cp libmetis/Release/metis.lib /c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/VC/Tools/MSVC/14.16.27023/lib/x64
+  cp ../include/metis.h /c/Program\ Files\ \(x86\)/Microsoft\ Visual\ Studio/2017/BuildTools/VC/Tools/MSVC/14.16.27023/include
+
+  cd ..
+fi
+
+cd ..

setup.py

@@ -16,10 +16,18 @@
 
 BUILD_DOCS = os.getenv('BUILD_DOCS', '0') == '1'
 
+WITH_METIS = False
+if os.getenv('WITH_METIS', '0') == '1':
+    WITH_METIS = True
+
 
 def get_extensions():
     Extension = CppExtension
     define_macros = []
+    libraries = []
+    if WITH_METIS:
+        define_macros += [('WITH_METIS', None)]
+        libraries += ['metis']
     extra_compile_args = {'cxx': []}
     extra_link_args = []
 
@@ -32,9 +40,9 @@ def get_extensions():
         extra_compile_args['nvcc'] = nvcc_flags
 
         if sys.platform == 'win32':
-            extra_link_args = ['cusparse.lib']
+            extra_link_args += ['cusparse.lib']
         else:
-            extra_link_args = ['-lcusparse', '-l', 'cusparse']
+            extra_link_args += ['-lcusparse', '-l', 'cusparse']
 
     extensions_dir = osp.join(osp.dirname(osp.abspath(__file__)), 'csrc')
     main_files = glob.glob(osp.join(extensions_dir, '*.cpp'))
@@ -59,6 +67,7 @@ def get_extensions():
             define_macros=define_macros,
             extra_compile_args=extra_compile_args,
             extra_link_args=extra_link_args,
+            libraries=libraries,
         )
         extensions += [extension]
 
@@ -71,7 +80,7 @@ def get_extensions():
 
 setup(
     name='torch_sparse',
-    version='0.5.1',
+    version='0.6.0',
     author='Matthias Fey',
     author_email='[email protected]',
     url='https://github.com/rusty1s/pytorch_sparse',

test/test_metis.py

@@ -0,0 +1,17 @@
+import pytest
+import torch
+from torch_sparse.tensor import SparseTensor
+
+from .utils import devices
+
+
+@pytest.mark.parametrize('device', devices)
+def test_metis(device):
+    mat = SparseTensor.from_dense(torch.randn((6, 6), device=device))
+    mat, partptr, perm = mat.partition(num_parts=2, recursive=False)
+    assert partptr.numel() == 3
+    assert perm.numel() == 6
+
+    mat, partptr, perm = mat.partition(num_parts=2, recursive=True)
+    assert partptr.numel() == 3
+    assert perm.numel() == 6

test/test_permute.py

@@ -0,0 +1,17 @@
+import pytest
+import torch
+from torch_sparse.tensor import SparseTensor
+
+from .utils import devices, tensor
+
+
+@pytest.mark.parametrize('device', devices)
+def test_permute(device):
+    row, col = tensor([[0, 0, 1, 2, 2], [0, 1, 0, 1, 2]], torch.long, device)
+    value = tensor([1, 2, 3, 4, 5], torch.float, device)
+    adj = SparseTensor(row=row, col=col, value=value)
+
+    row, col, value = adj.permute(torch.tensor([1, 0, 2])).coo()
+    assert row.tolist() == [0, 1, 1, 2, 2]
+    assert col.tolist() == [1, 0, 1, 0, 2]
+    assert value.tolist() == [3, 2, 1, 4, 5]

test/test_storage.py

@@ -93,8 +93,8 @@ def test_utility(dtype, device):
     storage = storage.set_value(value, layout='coo')
     assert storage.value().tolist() == [1, 2, 3, 4]
 
-    storage = storage.sparse_resize([3, 3])
-    assert storage.sparse_sizes() == [3, 3]
+    storage = storage.sparse_resize((3, 3))
+    assert storage.sparse_sizes() == (3, 3)
 
     new_storage = storage.copy()
     assert new_storage != storage

torch_sparse/__init__.py

@@ -3,11 +3,13 @@
 
 import torch
 
-__version__ = '0.5.1'
+__version__ = '0.6.0'
 expected_torch_version = (1, 4)
 
 try:
-    for library in ['_version', '_convert', '_diag', '_spmm', '_spspmm']:
+    for library in [
+            '_version', '_convert', '_diag', '_spmm', '_spspmm', '_metis'
+    ]:
         torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
             library, [osp.dirname(__file__)]).origin)
 except OSError as e:
@@ -45,12 +47,14 @@
 from .select import select  # noqa
 from .index_select import index_select, index_select_nnz  # noqa
 from .masked_select import masked_select, masked_select_nnz  # noqa
+from .permute import permute  # noqa
 from .diag import remove_diag, set_diag, fill_diag  # noqa
 from .add import add, add_, add_nnz, add_nnz_  # noqa
 from .mul import mul, mul_, mul_nnz, mul_nnz_  # noqa
 from .reduce import sum, mean, min, max  # noqa
 from .matmul import matmul  # noqa
 from .cat import cat, cat_diag  # noqa
+from .metis import partition  # noqa
 
 from .convert import to_torch_sparse, from_torch_sparse  # noqa
 from .convert import to_scipy, from_scipy  # noqa
@@ -71,6 +75,7 @@
     'index_select_nnz',
     'masked_select',
     'masked_select_nnz',
+    'permute',
     'remove_diag',
     'set_diag',
     'fill_diag',
@@ -89,6 +94,7 @@
     'matmul',
     'cat',
     'cat_diag',
+    'partition',
     'to_torch_sparse',
     'from_torch_sparse',
     'to_scipy',

torch_sparse/add.py

@@ -5,7 +5,6 @@
 from torch_sparse.tensor import SparseTensor
 
 
-@torch.jit.script
 def add(src: SparseTensor, other: torch.Tensor) -> SparseTensor:
     rowptr, col, value = src.csr()
     if other.size(0) == src.size(0) and other.size(1) == 1:  # Row-wise...
@@ -24,7 +23,6 @@ def add(src: SparseTensor, other: torch.Tensor) -> SparseTensor:
     return src.set_value(value, layout='coo')
 
 
-@torch.jit.script
 def add_(src: SparseTensor, other: torch.Tensor) -> SparseTensor:
     rowptr, col, value = src.csr()
     if other.size(0) == src.size(0) and other.size(1) == 1:  # Row-wise...
@@ -44,7 +42,6 @@ def add_(src: SparseTensor, other: torch.Tensor) -> SparseTensor:
     return src.set_value_(value, layout='coo')
 
 
-@torch.jit.script
 def add_nnz(src: SparseTensor, other: torch.Tensor,
             layout: Optional[str] = None) -> SparseTensor:
     value = src.storage.value()
@@ -55,7 +52,6 @@ def add_nnz(src: SparseTensor, other: torch.Tensor,
     return src.set_value(value, layout=layout)
 
 
-@torch.jit.script
 def add_nnz_(src: SparseTensor, other: torch.Tensor,
              layout: Optional[str] = None) -> SparseTensor:
     value = src.storage.value()

torch_sparse/cat.py

@@ -5,7 +5,6 @@
 from torch_sparse.tensor import SparseTensor
 
 
-@torch.jit.script
 def cat(tensors: List[SparseTensor], dim: int) -> SparseTensor:
     assert len(tensors) > 0
     if dim < 0:
@@ -142,7 +141,6 @@ def cat(tensors: List[SparseTensor], dim: int) -> SparseTensor:
             '[{-tensors[0].dim()}, {tensors[0].dim() - 1}], but got {dim}.')
 
 
-@torch.jit.script
 def cat_diag(tensors: List[SparseTensor]) -> SparseTensor:
     assert len(tensors) > 0
 

torch_sparse/coalesce.py

@@ -20,6 +20,6 @@ def coalesce(index, value, m, n, op="add"):
     """
 
     storage = SparseStorage(row=index[0], col=index[1], value=value,
-                            sparse_sizes=torch.Size([m, n]), is_sorted=False)
+                            sparse_sizes=(m, n), is_sorted=False)
     storage = storage.coalesce(reduce=op)
     return torch.stack([storage.row(), storage.col()], dim=0), storage.value()

torch_sparse/convert.py

@@ -5,7 +5,7 @@
 
 
 def to_torch_sparse(index, value, m, n):
-    return torch.sparse_coo_tensor(index.detach(), value, torch.Size([m, n]))
+    return torch.sparse_coo_tensor(index.detach(), value, (m, n))
 
 
 def from_torch_sparse(A):