dgl.graphbolt.unique_and_compact_csc_formats

dgl.graphbolt.unique_and_compact_csc_formats(csc_formats: Tuple[Tensor, Tensor] | Dict[str, Tuple[Tensor, Tensor]], unique_dst_nodes: Tensor | Dict[str, Tensor], rank: int = 0, world_size: int = 1, async_op: bool = False)[source]

Compact csc formats and return unique nodes (per type). The rank and world_size parameters are relevant when using Cooperative Minibatching, which was initially proposed in `Deep Graph Library PR#4337<https://github.com/dmlc/dgl/pull/4337>`__ and was later first fully described in Cooperative Minibatching in Graph Neural Networks. Cooperation between the GPUs eliminates duplicate work performed across the GPUs due to the overlapping sampled k-hop neighborhoods of seed nodes when performing GNN minibatching.

When world_size is greater than 1, then the given ids are partitioned between the available ranks. The ids corresponding to the given rank are guaranteed to come before the ids of other ranks. To do this, the partitioned ids are rotated backwards by the given rank so that the ids are ordered as: [rank, rank + 1, world_size, 0, …, rank - 1]. This is supported only for Volta and later generation NVIDIA GPUs.

Parameters:

  • csc_formats (Union[CSCFormatBase, Dict(str, CSCFormatBase)]) – CSC formats representing source-destination edges. - If csc_formats is a CSCFormatBase: It means the graph is homogeneous. Also, indptr and indice in it should be torch.tensor representing source and destination pairs in csc format. And IDs inside are homogeneous ids. - If csc_formats is a Dict[str, CSCFormatBase]: The keys should be edge type and the values should be csc format node pairs. And IDs inside are heterogeneous ids.

  • unique_dst_nodes (torch.Tensor or Dict[str, torch.Tensor]) – Unique nodes of all destination nodes in the node pairs. - If unique_dst_nodes is a tensor: It means the graph is homogeneous. - If csc_formats is a dictionary: The keys are node type and the values are corresponding nodes. And IDs inside are heterogeneous ids.

  • rank (int) – The rank of the current process.

  • world_size (int) – The number of processes.

  • async_op (bool) – Boolean indicating whether the call is asynchronous. If so, the result can be obtained by calling wait on the returned future.

Returns:

The compacted csc formats, where node IDs are replaced with mapped node IDs, and the unique nodes (per type). β€œCompacted csc formats” indicates that the node IDs in the input node pairs are replaced with mapped node IDs, where each type of node is mapped to a contiguous space of IDs ranging from 0 to N. The unique nodes offsets tensor partitions the unique_nodes tensor. Has size world_size + 1 and unique_nodes[offsets[i]: offsets[i + 1]] belongs to the rank (rank + i) % world_size.

Return type:

Tuple[unique_nodes, csc_formats, unique_nodes_offsets]

Examples

>>> import dgl.graphbolt as gb
>>> N1 = torch.LongTensor([1, 2, 2])
>>> N2 = torch.LongTensor([5, 5, 6])
>>> unique_dst = {
...     "n1": torch.LongTensor([1, 2]),
...     "n2": torch.LongTensor([5, 6])}
>>> csc_formats = {
...     "n1:e1:n2": gb.CSCFormatBase(indptr=torch.tensor([0, 2, 3]),indices=N1),
...     "n2:e2:n1": gb.CSCFormatBase(indptr=torch.tensor([0, 1, 3]),indices=N2)}
>>> unique_nodes, compacted_csc_formats, _ = gb.unique_and_compact_csc_formats(
...     csc_formats, unique_dst
... )
>>> print(unique_nodes)
{'n1': tensor([1, 2]), 'n2': tensor([5, 6])}
>>> print(compacted_csc_formats)
{"n1:e1:n2": CSCFormatBase(indptr=torch.tensor([0, 2, 3]),
                           indices=torch.tensor([0, 1, 1])),
 "n2:e2:n1": CSCFormatBase(indptr=torch.tensor([0, 1, 3]),
                           indices=torch.Longtensor([0, 0, 1]))}