Use optimized implementation in softmax operation (pyg-team#8399)

This PR uses optimized `softmax_csr` operation (introduced in [pyg-lib @
264](pyg-team/pyg-lib#264)), when given is a CPU
tensor, and softmax groups are defined via `ptr`.

CHANGELOG.md

@@ -27,6 +27,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Changed
 
+- Made `utils.softmax` faster via `softmax_csr` ([#8399](https://github.com/pyg-team/pytorch_geometric/pull/8399))
 - Made `utils.mask.mask_select` faster ([#8369](https://github.com/pyg-team/pytorch_geometric/pull/8369))
 - Update `DistNeighborSampler` for homogeneous graphs ([#8209](https://github.com/pyg-team/pytorch_geometric/pull/8209), [#8367](https://github.com/pyg-team/pytorch_geometric/pull/8367))
 - Update `GraphStore` and `FeatureStore` to support distributed training ([#8083](https://github.com/pyg-team/pytorch_geometric/pull/8083))

test/utils/test_softmax.py

@@ -5,6 +5,9 @@
 from torch_geometric.profile import benchmark
 from torch_geometric.utils import softmax
 
+CALCULATION_VIA_PTR_AVAILABLE = (torch_geometric.typing.WITH_SOFTMAX
+                                 or torch_geometric.typing.WITH_TORCH_SCATTER)
+
 
 def test_softmax():
     src = torch.tensor([1., 1., 1., 1.])
@@ -13,7 +16,7 @@ def test_softmax():
 
     out = softmax(src, index)
     assert out.tolist() == [0.5, 0.5, 1, 1]
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert softmax(src, None, ptr).tolist() == out.tolist()
     else:
         with pytest.raises(ImportError):
@@ -22,7 +25,7 @@ def test_softmax():
     src = src.view(-1, 1)
     out = softmax(src, index)
     assert out.tolist() == [[0.5], [0.5], [1], [1]]
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert softmax(src, None, ptr).tolist() == out.tolist()
 
     jit = torch.jit.script(softmax)
@@ -52,22 +55,22 @@ def test_softmax_dim():
 
     src = torch.randn(4)
     assert torch.allclose(softmax(src, index, dim=0), src.softmax(dim=0))
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert torch.allclose(softmax(src, ptr=ptr, dim=0), src.softmax(dim=0))
 
     src = torch.randn(4, 16)
     assert torch.allclose(softmax(src, index, dim=0), src.softmax(dim=0))
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert torch.allclose(softmax(src, ptr=ptr, dim=0), src.softmax(dim=0))
 
     src = torch.randn(4, 4)
     assert torch.allclose(softmax(src, index, dim=-1), src.softmax(dim=-1))
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert torch.allclose(softmax(src, ptr=ptr, dim=-1), src.softmax(-1))
 
     src = torch.randn(4, 4, 16)
     assert torch.allclose(softmax(src, index, dim=1), src.softmax(dim=1))
-    if torch_geometric.typing.WITH_TORCH_SCATTER:
+    if CALCULATION_VIA_PTR_AVAILABLE:
         assert torch.allclose(softmax(src, ptr=ptr, dim=1), src.softmax(dim=1))
 
 

torch_geometric/typing.py

@@ -40,6 +40,7 @@
             WITH_GMM = False
             WITH_SEGMM = False
     WITH_SAMPLED_OP = hasattr(pyg_lib.ops, 'sampled_add')
+    WITH_SOFTMAX = hasattr(pyg_lib.ops, 'softmax_csr')
     WITH_INDEX_SORT = hasattr(pyg_lib.ops, 'index_sort')
     WITH_METIS = hasattr(pyg_lib, 'partition')
     WITH_EDGE_TIME_NEIGHBOR_SAMPLE = ('edge_time' in inspect.signature(
@@ -55,6 +56,7 @@
     WITH_GMM = False
     WITH_SEGMM = False
     WITH_SAMPLED_OP = False
+    WITH_SOFTMAX = False
     WITH_INDEX_SORT = False
     WITH_METIS = False
     WITH_EDGE_TIME_NEIGHBOR_SAMPLE = False

torch_geometric/utils/softmax.py

@@ -2,6 +2,8 @@
 
 from torch import Tensor
 
+import torch_geometric.typing
+from torch_geometric.typing import pyg_lib
 from torch_geometric.utils import scatter, segment
 from torch_geometric.utils.num_nodes import maybe_num_nodes
 
@@ -50,6 +52,9 @@ def softmax(
                 [0.8062, 0.1938, 1.0000, 1.0000]])
     """
     if ptr is not None:
+        if (src.device.type == 'cpu'
+                and torch_geometric.typing.WITH_SOFTMAX):  # pragma: no cover
+            return pyg_lib.ops.softmax_csr(src, ptr, dim)
         dim = dim + src.dim() if dim < 0 else dim
         size = ([1] * dim) + [-1]
         count = ptr[1:] - ptr[:-1]