Merge branch 'master' of github.com:rusty1s/pytorch_sparse

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