Repository: pyg-team/pytorch_geometric
Pull request #8148 (2023-10-06 18:22:50)

Introduce global `torch_geometric.backend.use_segment_matmul` flag


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c1b850071..c3ac19f6b 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the option to override usage of `segment_matmul`/`grouped_matmul` via the `torch_geometric.backend.use_segment_matmul` flag ([#8148](https://github.com/pyg-team/pytorch_geometric/pull/8148))
 - Added support for PyTorch 2.1.0 ([#8134](https://github.com/pyg-team/pytorch_geometric/pull/8134))
 - Added the `NeuroGraphDataset` benchmark collection ([#8122](https://github.com/pyg-team/pytorch_geometric/pull/8122))
 - Added support for a node-level `mask` tensor in `dense_to_sparse` ([#8117](https://github.com/pyg-team/pytorch_geometric/pull/8117))
@@ -69,7 +70,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added the `HM` personalized fashion recommendation dataset ([#7515](https://github.com/pyg-team/pytorch_geometric/pull/7515))
 - Added the `GraphMixer` model ([#7501](https://github.com/pyg-team/pytorch_geometric/pull/7501), [#7459](https://github.com/pyg-team/pytorch_geometric/pull/7459))
 - Added the `disable_dynamic_shape` experimental flag ([#7246](https://github.com/pyg-team/pytorch_geometric/pull/7246), [#7534](https://github.com/pyg-team/pytorch_geometric/pull/7534))
-- Added the option to override `use_segmm` selection in `HeteroLinear` ([#7474](https://github.com/pyg-team/pytorch_geometric/pull/7474))
 - Added the `MovieLens-1M` heterogeneous dataset ([#7479](https://github.com/pyg-team/pytorch_geometric/pull/7479))
 - Added a CPU-based and GPU-based `map_index` implementation ([#7493](https://github.com/pyg-team/pytorch_geometric/pull/7493), [#7764](https://github.com/pyg-team/pytorch_geometric/pull/7764) [#7765](https://github.com/pyg-team/pytorch_geometric/pull/7765))
 - Added the `AmazonBook` heterogeneous dataset ([#7483](https://github.com/pyg-team/pytorch_geometric/pull/7483))
diff --git a/torch_geometric/backend.py b/torch_geometric/backend.py
new file mode 100644
index 000000000..3beb4a4ec
--- /dev/null
+++ b/torch_geometric/backend.py
@@ -0,0 +1,53 @@
+from typing import Optional
+
+import torch
+
+# If set to `True`, PyG is configured to use the `segment_matmul` and
+# `grouped_matmul` kernels from `pyg-lib` to parallelize matrix multiplication
+# across segments/groups of potentially varying size.
+# If set to `None`, will automatically decide whether to utilize
+# `segment_matmul` and `grouped_matmul` based on input sizes.
+# Requires `pyg-lib` to be installed.
+use_segment_matmul: Optional[bool] = None
+
+# Helper functions ############################################################
+
+
+def use_segment_matmul_heuristic(
+    num_segments: int,
+    max_segment_size: int,
+    in_channels: int,
+    out_channels: int,
+) -> bool:
+    r"""A heuristic based on input sizes to determine whether the usage of
+    :meth:`segment_matmul` can speed up computation."""
+    # NOTE This heuristic was learned on an A100 via sklearn using a simple
+    # StandardScaler() -> LinearSVC() model.
+    x = torch.tensor([
+        num_segments,
+        max_segment_size,
+        in_channels,
+        out_channels,
+    ])
+    mean = torch.tensor([
+        125.11603189,
+        12133.21523472,
+        163.81222321,
+        32.43755536,
+    ])
+    std = torch.tensor([
+        163.34480422,
+        27572.94543809,
+        177.6426489,
+        56.82103934,
+    ])
+    weight = torch.tensor([
+        2.43877659e+00,
+        1.67583047e+00,
+        -5.20527282e-04,
+        3.43925501e-01,
+    ])
+    bias = 1.20236999
+
+    x = (x - mean) / std
+    return bool(x @ weight >= bias)
diff --git a/torch_geometric/nn/conv/rgcn_conv.py b/torch_geometric/nn/conv/rgcn_conv.py
index e7ff9f342..016d5687f 100644
--- a/torch_geometric/nn/conv/rgcn_conv.py
+++ b/torch_geometric/nn/conv/rgcn_conv.py
@@ -5,6 +5,7 @@ from torch import Tensor
 from torch.nn import Parameter
 from torch.nn import Parameter as Param
 
+import torch_geometric.backend
 import torch_geometric.typing
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.inits import glorot, zeros
@@ -16,7 +17,6 @@ from torch_geometric.typing import (
     torch_sparse,
 )
 from torch_geometric.utils import index_sort, one_hot, scatter, spmm
-from torch_geometric.utils.hetero import segmatmul_heuristic
 from torch_geometric.utils.sparse import index2ptr
 
 
@@ -127,11 +127,13 @@ class RGCNConv(MessagePassing):
         self.num_bases = num_bases
         self.num_blocks = num_blocks
         self.is_sorted = is_sorted
-        self.use_segmm: int = -1
+
         if isinstance(in_channels, int):
             in_channels = (in_channels, in_channels)
         self.in_channels_l = in_channels[0]
 
+        self._use_segment_matmul_heuristic_output: Optional[bool] = None
+
         if num_bases is not None:
             self.weight = Parameter(
                 torch.empty(num_bases, in_channels[0], out_channels))
@@ -229,19 +231,37 @@ class RGCNConv(MessagePassing):
                 out = out + h.contiguous().view(-1, self.out_channels)
 
         else:  # No regularization/Basis-decomposition ========================
-            if (torch_geometric.typing.WITH_SEGMM and self.num_bases is None
-                    and x_l.is_floating_point() and isinstance(
-                        edge_index, Tensor)) and (self.use_segmm == -1
-                                                  or bool(self.use_segmm)):
+
+            use_segment_matmul = torch_geometric.backend.use_segment_matmul
+            # If `use_segment_matmul` is not specified, use a simple heuristic
+            # to determine whether `segment_matmul` can speed up computation
+            # given the observed input sizes:
+            if use_segment_matmul is None:
+                segment_count = scatter(torch.ones_like(edge_type), edge_type,
+                                        dim_size=self.num_relations)
+
+                self._use_segment_matmul_heuristic_output = (
+                    torch_geometric.backend.use_segment_matmul_heuristic(
+                        num_segments=self.num_relations,
+                        max_segment_size=int(segment_count.max()),
+                        in_channels=self.weight.size(1),
+                        out_channels=self.weight.size(2),
+                    ))
+
+                assert self._use_segment_matmul_heuristic_output is not None
+                use_segment_matmul = self._use_segment_matmul_heuristic_output
+
+            use_segment_matmul &= self.num_bases is None
+            use_segment_matmul &= x_l.is_floating_point()
+            use_segment_matmul &= isinstance(edge_index, Tensor)
+
+            if use_segment_matmul and torch_geometric.typing.WITH_SEGMM:
                 if not self.is_sorted:
                     if (edge_type[1:] < edge_type[:-1]).any():
                         edge_type, perm = index_sort(
                             edge_type, max_value=self.num_relations)
                         edge_index = edge_index[:, perm]
                 edge_type_ptr = index2ptr(edge_type, self.num_relations)
-                if self.use_segmm == -1:
-                    self.use_segmm = segmatmul_heuristic(
-                        x_l, edge_type_ptr, self.weight)
                 out = self.propagate(edge_index, x=x_l,
                                      edge_type_ptr=edge_type_ptr, size=size)
             else:
diff --git a/torch_geometric/nn/dense/linear.py b/torch_geometric/nn/dense/linear.py
index d7804c12d..7d7cd223e 100644
--- a/torch_geometric/nn/dense/linear.py
+++ b/torch_geometric/nn/dense/linear.py
@@ -4,21 +4,21 @@ from typing import Any, Dict, Optional, Union
 
 import torch
 import torch.nn.functional as F
-from torch import Tensor, nn
+from torch import Tensor
 from torch.nn.parameter import Parameter
 
+import torch_geometric.backend
 import torch_geometric.typing
 from torch_geometric.nn import inits
 from torch_geometric.typing import pyg_lib
-from torch_geometric.utils import index_sort
-from torch_geometric.utils.hetero import segmatmul_heuristic
+from torch_geometric.utils import index_sort, scatter
 from torch_geometric.utils.sparse import index2ptr
 
 
 def is_uninitialized_parameter(x: Any) -> bool:
-    if not hasattr(nn.parameter, 'UninitializedParameter'):
+    if not hasattr(torch.nn.parameter, 'UninitializedParameter'):
         return False
-    return isinstance(x, nn.parameter.UninitializedParameter)
+    return isinstance(x, torch.nn.parameter.UninitializedParameter)
 
 
 def reset_weight_(weight: Tensor, in_channels: int,
@@ -96,7 +96,7 @@ class Linear(torch.nn.Module):
         if in_channels > 0:
             self.weight = Parameter(torch.empty(out_channels, in_channels))
         else:
-            self.weight = nn.parameter.UninitializedParameter()
+            self.weight = torch.nn.parameter.UninitializedParameter()
             self._hook = self.register_forward_pre_hook(
                 self.initialize_parameters)
 
@@ -160,7 +160,7 @@ class Linear(torch.nn.Module):
 
         if weight is not None and is_uninitialized_parameter(weight):
             self.in_channels = -1
-            self.weight = nn.parameter.UninitializedParameter()
+            self.weight = torch.nn.parameter.UninitializedParameter()
             if not hasattr(self, '_hook'):
                 self._hook = self.register_forward_pre_hook(
                     self.initialize_parameters)
@@ -198,13 +198,6 @@ class HeteroLinear(torch.nn.Module):
             :obj:`type_vec` is sorted. This avoids internal re-sorting of the
             data and can improve runtime and memory efficiency.
             (default: :obj:`False`)
-        use_segmm (bool, optional): If set to :obj:`True` and :obj:`pyg-lib` is
-            installed, this module will use the fused :obj:`segment_matmul`
-            kernel to parallelize the linear transformation across types. If
-            set to :obj:`False`, :obj:`segment_matmul` will not be used. If
-            left as :obj:`None` and :obj:`pyg-lib` is installed, the module
-            will determine heuristically whether to use :obj:`segment_matmul`.
-            (default: :obj:`None`)
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.Linear`.
 
@@ -220,7 +213,6 @@ class HeteroLinear(torch.nn.Module):
         out_channels: int,
         num_types: int,
         is_sorted: bool = False,
-        use_segmm: Optional[bool] = None,
         **kwargs,
     ):
         super().__init__()
@@ -229,11 +221,12 @@ class HeteroLinear(torch.nn.Module):
         self.out_channels = out_channels
         self.num_types = num_types
         self.is_sorted = is_sorted
-        self.use_segmm: int = -1 if use_segmm is None else int(use_segmm)
         self.kwargs = kwargs
 
+        self._use_segment_matmul_heuristic_output: Optional[bool] = None
+
         if self.in_channels == -1:
-            self.weight = nn.parameter.UninitializedParameter()
+            self.weight = torch.nn.parameter.UninitializedParameter()
             self._hook = self.register_forward_pre_hook(
                 self.initialize_parameters)
         else:
@@ -258,8 +251,27 @@ class HeteroLinear(torch.nn.Module):
             x (torch.Tensor): The input features.
             type_vec (torch.Tensor): A vector that maps each entry to a type.
         """
-        if (torch_geometric.typing.WITH_SEGMM
-                and (self.use_segmm == -1 or bool(self.use_segmm))):
+        use_segment_matmul = torch_geometric.backend.use_segment_matmul
+        # If `use_segment_matmul` is not specified, use a simple heuristic to
+        # determine whether `segment_matmul` can speed up computation given the
+        # observed input sizes:
+        if use_segment_matmul is None:
+            if self._use_segment_matmul_heuristic_output is None:
+                segment_count = scatter(torch.ones_like(type_vec), type_vec,
+                                        dim_size=self.num_types, reduce='sum')
+
+                self._use_segment_matmul_heuristic_output = (
+                    torch_geometric.backend.use_segment_matmul_heuristic(
+                        num_segments=self.num_types,
+                        max_segment_size=int(segment_count.max()),
+                        in_channels=self.weight.size(1),
+                        out_channels=self.weight.size(2),
+                    ))
+
+            assert self._use_segment_matmul_heuristic_output is not None
+            use_segment_matmul = self._use_segment_matmul_heuristic_output
+
+        if use_segment_matmul and torch_geometric.typing.WITH_SEGMM:
             assert self.weight is not None
 
             perm: Optional[Tensor] = None
@@ -269,9 +281,6 @@ class HeteroLinear(torch.nn.Module):
                     x = x[perm]
 
             type_vec_ptr = index2ptr(type_vec, self.num_types)
-            if self.use_segmm == -1:
-                self.use_segmm = segmatmul_heuristic(x, type_vec_ptr,
-                                                     self.weight)
             out = pyg_lib.ops.segment_matmul(x, type_vec_ptr, self.weight)
             if self.bias is not None:
                 out += self.bias[type_vec]
@@ -328,13 +337,6 @@ class HeteroDictLinear(torch.nn.Module):
         out_channels (int): Size of each output sample.
         types (List[Any], optional): The keys of the input dictionary.
             (default: :obj:`None`)
-        use_segmm (bool, optional): If set to :obj:`True` and :obj:`pyg-lib` is
-            installed, this module will use the fused :obj:`segment_matmul`
-            kernel to parallelize the linear transformation across types. If
-            set to :obj:`False`, :obj:`segment_matmul` will not be used. If
-            left as :obj:`None` and :obj:`pyg-lib` is installed, the module
-            will determine heuristically whether to use :obj:`segment_matmul`.
-            (default: :obj:`None`)
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.Linear`.
     """
@@ -343,7 +345,6 @@ class HeteroDictLinear(torch.nn.Module):
         in_channels: Union[int, Dict[Any, int]],
         out_channels: int,
         types: Optional[Any] = None,
-        use_segmm: Optional[bool] = None,
         **kwargs,
     ):
         super().__init__()
@@ -373,7 +374,6 @@ class HeteroDictLinear(torch.nn.Module):
 
         self.in_channels = in_channels
         self.out_channels = out_channels
-        self.use_segmm = use_segmm
         self.kwargs = kwargs
 
         self.lins = torch.nn.ModuleDict({
@@ -400,15 +400,14 @@ class HeteroDictLinear(torch.nn.Module):
         """
         out_dict = {}
 
-        # Only apply fused kernel for more than 10 types, otherwise default
-        # back to sequential computation (which is faster for these cases).
-        if self.use_segmm is None:
-            use_segmm = len(x_dict) >= 10
-        else:
-            use_segmm = self.use_segmm
+        # Only apply fused kernel for more than 10 types, otherwise use
+        # sequential computation (which is generally faster for these cases).
+        use_segment_matmul = torch_geometric.backend.use_segment_matmul
+        if use_segment_matmul is None:
+            use_segment_matmul = len(x_dict) >= 10
 
-        if (torch_geometric.typing.WITH_GMM and not torch.jit.is_scripting()
-                and use_segmm):
+        if (use_segment_matmul and torch_geometric.typing.WITH_GMM
+                and not torch.jit.is_scripting()):
             xs, weights, biases = [], [], []
             for key, lin in self.lins.items():
                 if key in x_dict:
diff --git a/torch_geometric/utils/hetero.py b/torch_geometric/utils/hetero.py
index 7715a4042..1e9b8c217 100644
--- a/torch_geometric/utils/hetero.py
+++ b/torch_geometric/utils/hetero.py
@@ -9,115 +9,6 @@ from torch_geometric.utils import is_sparse, to_edge_index
 from torch_geometric.utils.num_nodes import maybe_num_nodes_dict
 
 
-def learn_sklearn_heuristic():
-    import os
-    import time
-
-    from torch_geometric.nn.dense import HeteroLinear, Linear
-    os.environ['NVIDIA_TF32_OVERRIDE'] = '0'
-    fused_times = {}
-    loop_times = {}
-    try:
-        for num_nodes_per_type in [10**2, 10**3, 10**4, 10**5]:
-            for out_feats in [2, 4, 8, 16, 32, 64, 128, 256]:
-                for n_feats in [4, 8, 16, 32, 64, 128, 256, 512]:
-                    for num_types in [4, 8, 16, 32, 64, 128, 256, 512]:
-                        try:
-                            if n_feats < out_feats:
-                                continue
-                            print("benchmarking", num_types, "types w/",
-                                  num_nodes_per_type, "nodes per type and",
-                                  n_feats, "input features and", out_feats,
-                                  "outuput feats")
-                            x_dict = {
-                                'v' + str(i):
-                                torch.randn(
-                                    (num_nodes_per_type, n_feats)).cuda()
-                                for i in range(num_types)
-                            }
-                            x = torch.cat(list(x_dict.values()), dim=0)
-                            node_type = torch.cat([
-                                (j * torch.ones(x_j.shape[0])).long()
-                                for j, x_j in enumerate(x_dict.values())
-                            ]).cuda()
-                            lin = Linear(n_feats, out_feats).cuda()
-                            heterolin = HeteroLinear(n_feats, out_feats,
-                                                     len(list(x_dict.keys())),
-                                                     True).cuda()
-                            for i in range(60):
-                                if i == 10:
-                                    since = time.time()
-                                heterolin(x=x, type_vec=node_type)
-                            key = (num_types, num_nodes_per_type, n_feats,
-                                   out_feats)
-                            fused_times[key] = ((time.time() - since) / 50.0)
-                            print("Avg time for fuse based=", fused_times[key])
-                            for i in range(60):
-                                if i == 10:
-                                    since = time.time()
-                                o = x.new_empty(x.size(0), out_feats)
-                                for j in range(num_types):
-                                    mask = j == node_type
-                                    o[mask] = lin(x[mask])
-                            loop_times[key] = ((time.time() - since) / 50.0)
-                            print("Avg time for for-loop=", loop_times[key])
-                        except:  # noqa
-                            continue
-    except:  # noqa
-        pass
-    import numpy as np
-    X = np.zeros((len(loop_times), 4))
-    y = np.zeros(len(loop_times))
-    for i, key in enumerate(loop_times.keys()):
-        X[i, :] = key
-        loop_time, fused_time = loop_times[key], fused_times[key]
-        y[i] = int(fused_time <= loop_time)
-    from sklearn.pipeline import make_pipeline
-    from sklearn.preprocessing import StandardScaler
-    from sklearn.svm import LinearSVC
-    scaler = StandardScaler()
-    svm = LinearSVC()
-    clf = make_pipeline(scaler, svm)
-    clf.fit(X, y)
-
-    print("scaler mean=", scaler.mean_)
-    print("scaler scale=", scaler.scale_)
-    print("svm weights=", svm.coef_)
-    print("svm bias=", svm.intercept_)
-    # results on A100:
-    # scaler mean=
-    #     [  125.11603189 12133.21523472   163.81222321    32.43755536]
-    # scaler scale=
-    #     [  163.34480422 27572.94543809   177.6426489     56.82103934]
-    # svm weights=
-    #     [[ 2.43877659e+00  1.67583047e+00 -5.20527282e-04  3.43925501e-01]]
-    # svm bias=
-    #     [1.20236999]
-
-
-def segmatmul_heuristic(inputs: Tensor, type_ptr, weight: Tensor):
-    num_types = len(type_ptr) - 1
-    max_num_nodes_per_types = (type_ptr[1:] - type_ptr[:-1]).max()
-    in_feat = inputs.size(1)
-    out_feat = weight.size(-1)
-    # this heuristic was learned with learn_sklearn_heuristic on an A100
-    x = torch.tensor([
-        int(num_types),
-        int(max_num_nodes_per_types),
-        int(in_feat),
-        int(out_feat)
-    ])
-    scale_mean = torch.tensor(
-        [125.11603189, 12133.21523472, 163.81222321, 32.43755536])
-    scale_scale = torch.tensor(
-        [163.34480422, 27572.94543809, 177.6426489, 56.82103934])
-    svm_weights = torch.tensor(
-        [2.43877659e+00, 1.67583047e+00, -5.20527282e-04, 3.43925501e-01])
-    bias = 1.20236999
-    x = (x - scale_mean) / scale_scale
-    return int(x.dot(svm_weights) >= bias)
-
-
 def group_hetero_graph(edge_index_dict, num_nodes_dict=None):
     num_nodes_dict = maybe_num_nodes_dict(edge_index_dict, num_nodes_dict)
 


Test patch:
diff --git a/test/nn/dense/test_linear.py b/test/nn/dense/test_linear.py
index f1189286e..5d2493ecb 100644
--- a/test/nn/dense/test_linear.py
+++ b/test/nn/dense/test_linear.py
@@ -8,6 +8,7 @@ from torch import Tensor
 from torch.nn import Linear as PTLinear
 from torch.nn.parameter import UninitializedParameter
 
+import torch_geometric.backend
 from torch_geometric.nn import HeteroDictLinear, HeteroLinear, Linear
 from torch_geometric.profile import benchmark
 from torch_geometric.testing import withCUDA, withPackage
@@ -140,19 +141,23 @@ def test_hetero_linear_initializer():
 
 
 @withCUDA
-@pytest.mark.parametrize('use_segmm', [True, False])
-def test_hetero_linear_amp(device, use_segmm):
+@pytest.mark.parametrize('use_segment_matmul', [None, True, False])
+def test_hetero_linear_amp(device, use_segment_matmul):
     warnings.filterwarnings('ignore', '.*but CUDA is not available.*')
 
+    old_state = torch_geometric.backend.use_segment_matmul
+    torch_geometric.backend.use_segment_matmul = use_segment_matmul
+
     x = torch.randn(3, 16, device=device)
     type_vec = torch.tensor([0, 1, 2], device=device)
 
     lin = HeteroLinear(16, 32, num_types=3).to(device)
-    lin.use_segmm = use_segmm
 
     with torch.cuda.amp.autocast():
         assert lin(x, type_vec).size() == (3, 32)
 
+    torch_geometric.backend.use_segment_matmul = old_state
+
 
 @withCUDA
 def test_lazy_hetero_linear(device):