From 24acde7a01a264bd9d1205fe96113457f28a7184 Mon Sep 17 00:00:00 2001 From: tobiaspk Date: Mon, 18 May 2026 18:08:49 -0400 Subject: [PATCH 1/3] Boolean to indices, chunked --- src/segger/data/tile_dataset.py | 26 ++++++++++++++++++++++---- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/src/segger/data/tile_dataset.py b/src/segger/data/tile_dataset.py index 921573f..1b7804d 100644 --- a/src/segger/data/tile_dataset.py +++ b/src/segger/data/tile_dataset.py @@ -9,7 +9,7 @@ from .partition import PartitionDataset from .tiling import Tiling - +_INT_MAX = 2**31 - 1 class TileFitDataset(PartitionDataset): """ Partitions a PyG graph based on a geometric tiling of its nodes. @@ -215,6 +215,23 @@ def __getitem__(self, idx: int) -> Data | HeteroData: geometry = self.tiling.tiles[idx] return self._subset(geometry) + def _chunked_nonzero(mask: torch.Tensor, chunk: int = 2 ** 30) -> torch.Tensor: + """Helper function to compute nonzero indices in chunks to avoid INT_MAX limit. (issue: https://github.com/dpeerlab/segger/issues/44)""" + + # pytorch supports only INT_MAX elements for subsetting. + if mask.numel() <= _INT_MAX or mask.device.type != "cuda": + return mask.nonzero(as_tuple=False).flatten() + + # split nonzero into chunks + parts = [] + for i, m in enumerate(mask.split(chunk)): + idx = m.nonzero(as_tuple=False).flatten() + if idx.numel(): + parts.append(idx + i * chunk) + + # re-assemble + return torch.cat(parts) + def _subset(self, bounds: shapely.Polygon) -> Data | HeteroData: """Slices all node attributes within bounds. @@ -229,12 +246,12 @@ def _subset(self, bounds: shapely.Polygon) -> Data | HeteroData: for node_type in self.data.node_types: pos: torch.Tensor = self.data[node_type]['pos'] # Row indices of masked elements inside tile w/ margin - subset[node_type] = ( + subset[node_type] = self._chunked_nonzero( (pos[:, 0] >= outer[0]) & (pos[:, 0] < outer[2]) & (pos[:, 1] >= outer[1]) & (pos[:, 1] < outer[3]) - ).nonzero().flatten() + ) p_mask[node_type] = ( (pos[subset[node_type], 0] >= inner[0]) & (pos[subset[node_type], 0] <= inner[2]) & @@ -253,7 +270,8 @@ def _subset(self, bounds: shapely.Polygon) -> Data | HeteroData: (pos[:, 0] < outer[2]) & (pos[:, 1] >= outer[1]) & (pos[:, 1] < outer[3]) - ).nonzero().flatten() + ) + subset = self._chunked_nonzero(subset) sample = self.data.subgraph(subset) sample['predict_mask'] = ( (pos[subset, 0] >= inner[0]) & From ece355c6331a78368b11927675672975347ab6ad Mon Sep 17 00:00:00 2001 From: tobiaspk Date: Mon, 18 May 2026 18:09:43 -0400 Subject: [PATCH 2/3] Remove unused parameter --- src/segger/data/tile_dataset.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/segger/data/tile_dataset.py b/src/segger/data/tile_dataset.py index 1b7804d..601cfbc 100644 --- a/src/segger/data/tile_dataset.py +++ b/src/segger/data/tile_dataset.py @@ -260,7 +260,6 @@ def _subset(self, bounds: shapely.Polygon) -> Data | HeteroData: ) sample = self.data.subgraph(subset) sample.set_value_dict('predict_mask', p_mask) - sample.set_value_dict('global_index', subset) return sample else: # is homogenous Data @@ -279,7 +278,6 @@ def _subset(self, bounds: shapely.Polygon) -> Data | HeteroData: (pos[subset, 1] >= inner[1]) & (pos[subset, 1] <= inner[3]) ) - sample['global_index'] = subset return sample From a134e17e0248502a410532d095bbe8d22795ae00 Mon Sep 17 00:00:00 2001 From: tobiaspk Date: Tue, 19 May 2026 11:43:08 -0400 Subject: [PATCH 3/3] Implement chunked subgraphing as fallback. --- src/segger/__init__.py | 5 +++ src/segger/_patches.py | 74 +++++++++++++++++++++++++++++++++ src/segger/data/tile_dataset.py | 23 ++-------- 3 files changed, 82 insertions(+), 20 deletions(-) create mode 100644 src/segger/_patches.py diff --git a/src/segger/__init__.py b/src/segger/__init__.py index f1914d1..2f6eadc 100644 --- a/src/segger/__init__.py +++ b/src/segger/__init__.py @@ -13,6 +13,11 @@ torch.cuda.memory.change_current_allocator(rmm_torch_allocator) enable_statistics() +# Apply pytorch patches for issue pytorch/pytorch#51871 (CUDA nonzero INT_MAX limit). +# Must run BEFORE any segger module imports HeteroData / bipartite_subgraph. +from ._patches import apply as _apply_patches +_apply_patches() + def free_mem_str() -> str: stats = get_statistics() return ( diff --git a/src/segger/_patches.py b/src/segger/_patches.py new file mode 100644 index 0000000..51360b0 --- /dev/null +++ b/src/segger/_patches.py @@ -0,0 +1,74 @@ +"""Workaround for pytorch/pytorch#51871 (CUDA `nonzero` INT_MAX limit). + +Patches `torch_geometric.utils.bipartite_subgraph` and the references already +imported by `torch_geometric.data.hetero_data` / `._subgraph` so that +`HeteroData.subgraph` falls back to a chunked-nonzero path when the edge +tensor on CUDA exceeds INT_MAX (~2.15B) elements. + +See: https://github.com/dpeerlab/segger/issues/44 +""" +import torch +import torch_geometric.utils._subgraph as _sg +import torch_geometric.utils as _tgu +import torch_geometric.data.hetero_data as _hd +from torch_geometric.utils import index_to_mask +from torch_geometric.utils.map import map_index + +_INT_MAX = 2**31 - 1 +_pyg_bipartite = _sg.bipartite_subgraph + + +def chunked_nonzero(mask: torch.Tensor, chunk: int = 2**30) -> torch.Tensor: + """Chunked version of `mask.nonzero()` that works on CUDA tensors with > INT_MAX elements.""" + if mask.numel() <= _INT_MAX or mask.device.type != "cuda": + return mask.nonzero(as_tuple=False).flatten() + parts = [] + for i, m in enumerate(mask.split(chunk)): + idx = m.nonzero(as_tuple=False).flatten() + if idx.numel(): + parts.append(idx + i * chunk) + return torch.cat(parts) + + +def bipartite_safe(subset, edge_index, edge_attr=None, relabel_nodes=False, + size=None, return_edge_mask=False): + """ + Replacement for `torch_geometric.utils.bipartite_subgraph`. + Falls back to a chunked subgraph version when the edge_index is too large for CUDA. + """ + # original + if edge_index.numel() <= _INT_MAX or edge_index.device.type != "cuda": + return _pyg_bipartite(subset, edge_index, edge_attr, relabel_nodes, + size, return_edge_mask) + + # same as source + src_sub, dst_sub = subset + src_mask = index_to_mask(src_sub, size=size[0]) + dst_mask = index_to_mask(dst_sub, size=size[1]) + edge_mask = src_mask[edge_index[0]] & dst_mask[edge_index[1]] + + # replaced this + idx = chunked_nonzero(edge_mask) + + # same as source (but indices instead of mask) + edge_index = edge_index[:, idx] + edge_attr = edge_attr[edge_mask] if edge_attr is not None else None + if relabel_nodes: + src_index, _ = map_index(edge_index[0], src_sub, max_index=size[0], inclusive=True) + dst_index, _ = map_index(edge_index[1], dst_sub, max_index=size[1], inclusive=True) + edge_index = torch.stack([src_index, dst_index], dim=0) + return (edge_index, edge_attr, edge_mask) if return_edge_mask else (edge_index, edge_attr) + + +_patches_applied = False + + +def apply(): + """Apply the patches.""" + global _patches_applied + if _patches_applied: + return + _sg.bipartite_subgraph = bipartite_safe + _tgu.bipartite_subgraph = bipartite_safe + _hd.bipartite_subgraph = bipartite_safe + _patches_applied = True diff --git a/src/segger/data/tile_dataset.py b/src/segger/data/tile_dataset.py index 601cfbc..470c2b4 100644 --- a/src/segger/data/tile_dataset.py +++ b/src/segger/data/tile_dataset.py @@ -8,8 +8,8 @@ from .partition import PartitionDataset from .tiling import Tiling +from .._patches import chunked_nonzero as _chunked_nonzero -_INT_MAX = 2**31 - 1 class TileFitDataset(PartitionDataset): """ Partitions a PyG graph based on a geometric tiling of its nodes. @@ -215,23 +215,6 @@ def __getitem__(self, idx: int) -> Data | HeteroData: geometry = self.tiling.tiles[idx] return self._subset(geometry) - def _chunked_nonzero(mask: torch.Tensor, chunk: int = 2 ** 30) -> torch.Tensor: - """Helper function to compute nonzero indices in chunks to avoid INT_MAX limit. (issue: https://github.com/dpeerlab/segger/issues/44)""" - - # pytorch supports only INT_MAX elements for subsetting. - if mask.numel() <= _INT_MAX or mask.device.type != "cuda": - return mask.nonzero(as_tuple=False).flatten() - - # split nonzero into chunks - parts = [] - for i, m in enumerate(mask.split(chunk)): - idx = m.nonzero(as_tuple=False).flatten() - if idx.numel(): - parts.append(idx + i * chunk) - - # re-assemble - return torch.cat(parts) - def _subset(self, bounds: shapely.Polygon) -> Data | HeteroData: """Slices all node attributes within bounds. @@ -246,7 +229,7 @@ def _subset(self, bounds: shapely.Polygon) -> Data | HeteroData: for node_type in self.data.node_types: pos: torch.Tensor = self.data[node_type]['pos'] # Row indices of masked elements inside tile w/ margin - subset[node_type] = self._chunked_nonzero( + subset[node_type] = _chunked_nonzero( (pos[:, 0] >= outer[0]) & (pos[:, 0] < outer[2]) & (pos[:, 1] >= outer[1]) & @@ -270,7 +253,7 @@ def _subset(self, bounds: shapely.Polygon) -> Data | HeteroData: (pos[:, 1] >= outer[1]) & (pos[:, 1] < outer[3]) ) - subset = self._chunked_nonzero(subset) + subset = _chunked_nonzero(subset) sample = self.data.subgraph(subset) sample['predict_mask'] = ( (pos[subset, 0] >= inner[0]) &