Repository: pyg-team/pytorch_geometric
Pull request #8490 (2023-11-30 10:33:20)

Faster sparse-dense matrix multiplication in `EdgeIndex`


Patch:
diff --git a/torch_geometric/data/edge_index.py b/torch_geometric/data/edge_index.py
index ca9cbb4c0..cea6d6641 100644
--- a/torch_geometric/data/edge_index.py
+++ b/torch_geometric/data/edge_index.py
@@ -378,8 +378,8 @@ def apply_(
     out = out.as_subclass(EdgeIndex)
 
     # Copy metadata:
-    out._sparse_size = out.sparse_size
-    out._sort_order = out._sort_order
+    out._sparse_size = tensor.sparse_size
+    out._sort_order = tensor._sort_order
 
     # Convert cache:
     if tensor._rowptr is not None:
@@ -394,6 +394,12 @@ def apply_(
     return out
 
 
+@implements(torch.clone)
+@implements(Tensor.clone)
+def clone(tensor: EdgeIndex) -> EdgeIndex:
+    return apply_(tensor, Tensor.clone)
+
+
 @implements(Tensor.to)
 def to(tensor: EdgeIndex, *args, **kwargs) -> EdgeIndex:
     out = apply_(tensor, Tensor.to, *args, **kwargs)
@@ -690,6 +696,77 @@ else:
         return to_sparse_coo(tensor, value)
 
 
+ReduceType = Literal['sum']
+
+
+class SparseDenseMatmul(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        input: EdgeIndex,
+        other: Tensor,
+        input_value: Optional[Tensor] = None,
+        reduce: ReduceType = 'sum',
+    ) -> Tensor:
+
+        assert reduce in ['sum']
+
+        if other.requires_grad:
+            ctx.save_for_backward(input, input_value)
+
+        if input_value is None:
+            input_value = _get_value(
+                input.size(1),
+                dtype=other.dtype,
+                device=input.device,
+            )
+
+        adj = to_sparse_csr(input, input_value)  # Requires `sort_order="row"`.
+
+        return adj @ other
+
+    @staticmethod
+    def backward(
+        ctx,
+        out_grad: Tensor,
+    ) -> Tuple[None, Optional[Tensor], None, None]:
+
+        other_grad: Optional[Tensor] = None
+        if ctx.needs_input_grad[1]:
+            input, input_value = ctx.saved_tensors
+
+            # We need to compute adj.t() @ out_grad. To perform the transpose,
+            # we first sort by column and then create a CSR matrix from it.
+            # Note that the sort result is cached across multiple applications.
+            input, perm = input.sort_by(SortOrder.COL)
+            if input_value is not None:
+                input_value = input_value[perm]
+            else:
+                input_value = _get_value(
+                    input.size(1),
+                    dtype=out_grad.dtype,
+                    device=input.device,
+                )
+
+            # We can call `input.flip(0)` here and call `to_sparse_csr`, but
+            # creating the sparse CSR matrix from `colptr` directly is a bit
+            # more efficient:
+            adj_t = torch.sparse_csr_tensor(
+                crow_indices=input.get_colptr(),
+                col_indices=input[0],
+                values=input_value,
+                size=input.get_sparse_size()[::-1],
+                device=input.device,
+            )
+
+            other_grad = adj_t @ out_grad
+
+        if ctx.needs_input_grad[2]:
+            raise NotImplementedError
+
+        return None, other_grad, None, None
+
+
 @implements(torch.matmul)
 @implements(Tensor.matmul)
 def matmul(
@@ -697,12 +774,16 @@ def matmul(
     other: Union[Tensor, EdgeIndex],
     input_value: Optional[Tensor] = None,
     other_value: Optional[Tensor] = None,
-    reduce: Literal['sum'] = 'sum',
+    reduce: ReduceType = 'sum',
 ) -> Union[Tensor, Tuple[EdgeIndex, Tensor]]:
 
     assert reduce in ['sum']
 
-    # TODO Utilize available `CSC` representation for faster backward passes.
+    if not isinstance(other, EdgeIndex):
+        if other_value is not None:
+            raise ValueError("'other_value' not supported for sparse-dense "
+                             "matrix multiplication")
+        return SparseDenseMatmul.apply(input, other, input_value, reduce)
 
     if input._sort_order == SortOrder.COL:
         input = to_sparse_csc(input, input_value)
@@ -715,8 +796,4 @@ def matmul(
         else:
             other = to_sparse_csr(other, other_value)
 
-    elif other_value is not None:
-        raise ValueError("'other_value' not supported for sparse-dense "
-                         "matrix multiplication")
-
     return torch.matmul(input, other)


Test patch:
diff --git a/test/data/test_edge_index.py b/test/data/test_edge_index.py
index 141b6696b..cec284e53 100644
--- a/test/data/test_edge_index.py
+++ b/test/data/test_edge_index.py
@@ -42,6 +42,18 @@ def test_fill_cache():
     assert torch.equal(adj._colptr, torch.tensor([0, 1, 3, 4]))
 
 
+def test_clone():
+    adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
+
+    out = adj.clone()
+    assert isinstance(out, EdgeIndex)
+    assert out.sort_order == 'row'
+
+    out = torch.clone(adj)
+    assert isinstance(out, EdgeIndex)
+    assert out.sort_order == 'row'
+
+
 @withCUDA
 def test_to(device):
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]])
@@ -299,7 +311,7 @@ def test_to_sparse_csc():
     assert torch.equal(adj[0], out.row_indices())
 
 
-def test_matmul():
+def test_matmul_forward():
     x = torch.randn(3, 1)
     adj1 = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
     adj2 = EdgeIndex([[1, 0, 2, 1], [0, 1, 1, 2]], sort_order='col')
@@ -307,9 +319,6 @@ def test_matmul():
     out = adj1 @ x
     assert torch.allclose(out, adj1.to_dense() @ x)
 
-    out = adj2 @ x
-    assert torch.allclose(out, adj2.to_dense() @ x)
-
     out = adj1 @ adj1
     assert torch.allclose(out.to_dense(), adj1.to_dense() @ adj1.to_dense())
 
@@ -327,30 +336,28 @@ def test_matmul_input_value():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
 
     x = torch.randn(3, 1)
-    value = torch.randn(4)
+    input_value = torch.randn(4)
 
-    out = matmul(adj, x, input_value=value)
-    assert torch.allclose(out, to_dense(adj, value=value) @ x)
+    out = matmul(adj, x, input_value)
+    assert torch.allclose(out, to_dense(adj, value=input_value) @ x)
 
 
-def test_matmul_grad():
+def test_matmul_backward():
     adj = EdgeIndex([[0, 1, 1, 2], [1, 0, 2, 1]], sort_order='row')
 
     x1 = torch.randn(3, 1, requires_grad=True)
-    value = torch.randn(4, requires_grad=True)
+    input_value = torch.randn(4)
 
-    out = matmul(adj, x1, input_value=value)
+    out = matmul(adj, x1, input_value)
     grad_out = torch.randn_like(out)
     out.backward(grad_out)
 
     x2 = x1.detach().requires_grad_()
-    dense_adj = to_dense(adj, value=value).detach().requires_grad_()
+    dense_adj = to_dense(adj, value=input_value)
     out = dense_adj @ x2
     out.backward(grad_out)
 
     assert torch.allclose(x1.grad, x2.grad)
-    if torch_geometric.typing.WITH_PT21:  # TODO Investigate.
-        assert torch.allclose(value.grad, dense_adj.grad[adj[0], adj[1]])
 
 
 def test_save_and_load(tmp_path):