Enable index_sort (#306)

* Enable `index_sort`

* update

* update

---------

Co-authored-by: rusty1s <[email protected]>

Author: DamianSzwichtenberg

torch_sparse/storage.py

@@ -1,9 +1,10 @@
 import warnings
-from typing import Optional, List, Tuple
+from typing import List, Optional, Tuple
 
 import torch
-from torch_scatter import segment_csr, scatter_add
-from torch_sparse.utils import Final
+from torch_scatter import scatter_add, segment_csr
+
+from torch_sparse.utils import Final, index_sort
 
 layouts: Final[List[str]] = ['coo', 'csr', 'csc']
 
@@ -151,7 +152,8 @@ def __init__(
             idx[1:] *= self._sparse_sizes[1]
             idx[1:] += self._col
             if (idx[1:] < idx[:-1]).any():
-                perm = idx[1:].argsort()
+                max_value = self._sparse_sizes[0] * self._sparse_sizes[1]
+                _, perm = index_sort(idx[1:], max_value)
                 self._row = self.row()[perm]
                 self._col = self._col[perm]
                 if value is not None:
@@ -163,10 +165,20 @@ def __init__(
     def empty(self):
         row = torch.tensor([], dtype=torch.long)
         col = torch.tensor([], dtype=torch.long)
-        return SparseStorage(row=row, rowptr=None, col=col, value=None,
-                             sparse_sizes=(0, 0), rowcount=None, colptr=None,
-                             colcount=None, csr2csc=None, csc2csr=None,
-                             is_sorted=True, trust_data=True)
+        return SparseStorage(
+            row=row,
+            rowptr=None,
+            col=col,
+            value=None,
+            sparse_sizes=(0, 0),
+            rowcount=None,
+            colptr=None,
+            colcount=None,
+            csr2csc=None,
+            csc2csr=None,
+            is_sorted=True,
+            trust_data=True,
+        )
 
     def has_row(self) -> bool:
         return self._row is not None
@@ -209,8 +221,11 @@ def has_value(self) -> bool:
     def value(self) -> Optional[torch.Tensor]:
         return self._value
 
-    def set_value_(self, value: Optional[torch.Tensor],
-                   layout: Optional[str] = None):
+    def set_value_(
+        self,
+        value: Optional[torch.Tensor],
+        layout: Optional[str] = None,
+    ):
         if value is not None:
             if get_layout(layout) == 'csc':
                 value = value[self.csc2csr()]
@@ -221,8 +236,11 @@ def set_value_(self, value: Optional[torch.Tensor],
         self._value = value
         return self
 
-    def set_value(self, value: Optional[torch.Tensor],
-                  layout: Optional[str] = None):
+    def set_value(
+        self,
+        value: Optional[torch.Tensor],
+        layout: Optional[str] = None,
+    ):
         if value is not None:
             if get_layout(layout) == 'csc':
                 value = value[self.csc2csr()]
@@ -375,8 +393,11 @@ def colcount(self) -> torch.Tensor:
         if colptr is not None:
             colcount = colptr[1:] - colptr[:-1]
         else:
-            colcount = scatter_add(torch.ones_like(self._col), self._col,
-                                   dim_size=self._sparse_sizes[1])
+            colcount = scatter_add(
+                torch.ones_like(self._col),
+                self._col,
+                dim_size=self._sparse_sizes[1],
+            )
         self._colcount = colcount
         return colcount
 
@@ -389,7 +410,8 @@ def csr2csc(self) -> torch.Tensor:
             return csr2csc
 
         idx = self._sparse_sizes[0] * self._col + self.row()
-        csr2csc = idx.argsort()
+        max_value = self._sparse_sizes[0] * self._sparse_sizes[1]
+        _, csr2csc = index_sort(idx, max_value)
         self._csr2csc = csr2csc
         return csr2csc
 
@@ -401,7 +423,8 @@ def csc2csr(self) -> torch.Tensor:
         if csc2csr is not None:
             return csc2csr
 
-        csc2csr = self.csr2csc().argsort()
+        max_value = self._sparse_sizes[0] * self._sparse_sizes[1]
+        _, csc2csr = index_sort(self.csr2csc(), max_value)
         self._csc2csr = csc2csr
         return csc2csr
 
@@ -543,7 +566,8 @@ def type(self, dtype: torch.dtype, non_blocking: bool = False):
             else:
                 return self.set_value(
                     value.to(dtype=dtype, non_blocking=non_blocking),
-                    layout='coo')
+                    layout='coo',
+                )
         else:
             return self
 

torch_sparse/typing.py

@@ -0,0 +1,8 @@
+try:
+    import pyg_lib  # noqa
+    WITH_PYG_LIB = True
+    WITH_INDEX_SORT = hasattr(pyg_lib.ops, 'index_sort')
+except ImportError:
+    pyg_lib = object
+    WITH_PYG_LIB = False
+    WITH_INDEX_SORT = False

torch_sparse/utils.py

@@ -1,10 +1,25 @@
-from typing import Any
+from typing import Any, Optional, Tuple
+
+import torch
+
+import torch_sparse.typing
+from torch_sparse.typing import pyg_lib
 
 try:
     from typing_extensions import Final  # noqa
 except ImportError:
     from torch.jit import Final  # noqa
 
 
+def index_sort(
+        inputs: torch.Tensor,
+        max_value: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
+    r"""See pyg-lib documentation for more details:
+    https://pyg-lib.readthedocs.io/en/latest/modules/ops.html"""
+    if not torch_sparse.typing.WITH_INDEX_SORT:  # pragma: no cover
+        return inputs.sort()
+    return pyg_lib.ops.index_sort(inputs, max_value)
+
+
 def is_scalar(other: Any) -> bool:
     return isinstance(other, int) or isinstance(other, float)