physics_mcp/source/physicsnemo/models/gnn_layers/distributed_graph.py

# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
# SPDX-FileCopyrightText: All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
from dataclasses import dataclass
from typing import List, Optional

import torch
import torch.distributed as dist

from physicsnemo.distributed import (
    DistributedManager,
    all_gather_v,
    gather_v,
    indexed_all_to_all_v,
    scatter_v,
)

logger = logging.getLogger(__name__)


@dataclass
class GraphPartition:
    """
    Class acting as an "utility" structure to hold all relevant buffers and variables
    to define a graph partition and facilitate exchange of necessary buffers for
    message passing on a distributed graph.

    A global graph is assumed to be defined through a global CSC structure
    defining edges between source nodes and destination nodes which are assumed
    to be numbered indexed by contiguous IDs. Hence, features associated to both
    nodes and edges can be represented through dense feature tables globally.
    When partitioning graph and features, we distribute destination nodes and all
    their incoming edges on all ranks within the partition group based on a specified
    mapping. Based on this scheme, there will a be a difference between
    partitioned source nodes (partitioned features) and local source node
    IDs which refer to the node IDs within the local graph defined by the
    destination nodes on each rank. To allow message passing, communication
    primitives have to ensure to gather all corresponding features for all
    local source nodes based on the applied partitioning scheme. This also
    leads to the distinction of local source node IDs and remote source node
    IDs on each rank where the latter simply refers to the local row ID within
    the dense partitioning of node features and the former indicates the source
    of a message for each edge within each local graph.

    Parameters
    ----------
    partition_size : int
        size of partition
    partition_rank : int
        local rank of this partition w.r.t. group of partitions
    device : torch.device
        device handle for buffers within this partition rank
    """

    partition_size: int
    partition_rank: int
    device: torch.device
    # flag to indicate using adj matrix 1-D row-decomp
    matrix_decomp: bool = False

    # data structures defining partition
    # set in after initialization or during execution
    # of desired partition scheme

    # local CSR offsets defining local graph on each `partition_rank`
    local_offsets: Optional[torch.Tensor] = None
    # local CSR indices defining local graph on each `partition_rank`
    local_indices: Optional[torch.Tensor] = None
    # number of source nodes in local graph on each `partition_rank`
    num_local_src_nodes: int = -1
    # number of destination nodes in local graph on each `partition_rank`
    num_local_dst_nodes: int = -1
    # number of edges in local graph on each `partition_rank`
    num_local_indices: int = -1
    # mapping from global to local ID space (source node IDs)
    map_partitioned_src_ids_to_global: Optional[torch.Tensor] = None
    map_concatenated_local_src_ids_to_global: Optional[torch.Tensor] = None
    # mapping from local to global ID space (destination node IDs)
    map_partitioned_dst_ids_to_global: Optional[torch.Tensor] = None
    map_concatenated_local_dst_ids_to_global: Optional[torch.Tensor] = None
    # mapping from local to global ID space (edge IDs)
    map_partitioned_edge_ids_to_global: Optional[torch.Tensor] = None
    map_concatenated_local_edge_ids_to_global: Optional[torch.Tensor] = None
    # reverse mappings
    map_global_src_ids_to_concatenated_local: Optional[torch.Tensor] = None
    map_global_dst_ids_to_concatenated_local: Optional[torch.Tensor] = None
    map_global_edge_ids_to_concatenated_local: Optional[torch.Tensor] = None

    # utility lists and sizes required for exchange of messages
    # between graph partitions through distributed communication primitives

    # number of IDs each rank potentially sends to all other ranks
    sizes: Optional[List[List[int]]] = None
    # local indices of IDs current rank sends to all other ranks
    scatter_indices: Optional[List[torch.Tensor]] = None
    # number of global source nodes for each `partition_rank`
    num_src_nodes_in_each_partition: Optional[List[int]] = None
    # number of global destination nodes for each `partition_rank`
    num_dst_nodes_in_each_partition: Optional[List[int]] = None
    # number of global indices for each `partition_rank`
    num_indices_in_each_partition: Optional[List[int]] = None

    def __post_init__(self):
        # after partition_size has been set in __init__
        if self.partition_size <= 0:
            raise ValueError(f"Expected partition_size > 0, got {self.partition_size}")
        if not (0 <= self.partition_rank < self.partition_size):
            raise ValueError(
                f"Expected 0 <= partition_rank < {self.partition_size}, got {self.partiton_rank}"
            )

        if self.sizes is None:
            self.sizes = [
                [None for _ in range(self.partition_size)]
                for _ in range(self.partition_size)
            ]

        if self.scatter_indices is None:
            self.scatter_indices = [None] * self.partition_size
        if self.num_src_nodes_in_each_partition is None:
            self.num_src_nodes_in_each_partition = [None] * self.partition_size
        if self.num_dst_nodes_in_each_partition is None:
            self.num_dst_nodes_in_each_partition = [None] * self.partition_size
        if self.num_indices_in_each_partition is None:
            self.num_indices_in_each_partition = [None] * self.partition_size

    def to(self, *args, **kwargs):
        # move all tensors
        for attr in dir(self):
            attr_val = getattr(self, attr)
            if isinstance(attr_val, torch.Tensor):
                setattr(self, attr, attr_val.to(*args, **kwargs))

        # handle scatter_indices separately
        self.scatter_indices = [idx.to(*args, **kwargs) for idx in self.scatter_indices]

        return self


def partition_graph_with_id_mapping(
    global_offsets: torch.Tensor,
    global_indices: torch.Tensor,
    mapping_src_ids_to_ranks: torch.Tensor,
    mapping_dst_ids_to_ranks: torch.Tensor,
    partition_size: int,
    partition_rank: int,
    device: torch.device,
) -> GraphPartition:
    """
    Utility function which partitions a global graph given as CSC structure.
    It partitions both the global ID spaces for source nodes and destination nodes
    based on the corresponding mappings as well as the graph structure and edge IDs.
    For more details on partitioning in general see `GraphPartition`.
    The function performs the partitioning based on a global graph in CPU
    memory for each rank independently. It could be rewritten to e.g. only
    do it one rank and exchange the partitions or to an algorithm that also
    assumes an already distributed global graph, however, we expect global
    graphs to fit in CPU memory. After the partitioning, we can get rid off
    the larger one in CPU memory, only keep the local graphs on each GPU, and
    avoid tedious gather/scatter routines for exchanging partitions in the process.
    Note: It is up to the user to ensure that the provided mapping is valid. In particular,
    we expect each rank to receive a non-empty partition of node IDs.

    Parameters
    ----------
    global_offsets : torch.Tensor
        CSC offsets, can live on the CPU
    global_indices : torch.Tensor
        CSC indices, can live on the CPU
    mapping_src_ids_to_ranks: torch.Tensor
        maps each global ID from every source node to its partition rank
    mapping_dst_ids_to_ranks: torch.Tensor
        maps each global ID from every destination node to its partition rank
    partition_size : int
        number of process groups across which graph is partitioned,
        i.e. the number of graph partitions
    partition_rank : int
        rank within process group managing the distributed graph, i.e.
        the rank determining which partition the corresponding local rank
        will manage
    device : torch.device
        device connected to the passed partition rank, i.e. the device
        on which the local graph and related buffers will live on
    """

    # initialize graph partition
    graph_partition = GraphPartition(
        partition_size=partition_size, partition_rank=partition_rank, device=device
    )

    # --------------------------------------------------------------
    # initialize temporary variables used in computing the partition

    # global IDs of in each partition
    dst_nodes_in_each_partition = [None] * partition_size
    src_nodes_in_each_partition = [None] * partition_size
    num_dst_nodes_in_each_partition = [None] * partition_size
    num_src_nodes_in_each_partition = [None] * partition_size

    dtype = global_indices.dtype
    input_device = global_indices.device

    graph_partition.map_concatenated_local_src_ids_to_global = torch.empty_like(
        mapping_src_ids_to_ranks
    )
    graph_partition.map_concatenated_local_dst_ids_to_global = torch.empty_like(
        mapping_dst_ids_to_ranks
    )
    graph_partition.map_concatenated_local_edge_ids_to_global = torch.empty_like(
        global_indices
    )
    graph_partition.map_global_src_ids_to_concatenated_local = torch.empty_like(
        mapping_src_ids_to_ranks
    )
    graph_partition.map_global_dst_ids_to_concatenated_local = torch.empty_like(
        mapping_dst_ids_to_ranks
    )
    graph_partition.map_global_edge_ids_to_concatenated_local = torch.empty_like(
        global_indices
    )
    _map_global_src_ids_to_local = torch.empty_like(mapping_src_ids_to_ranks)

    # temporarily track cum-sum of nodes per partition for "concatenated_local_ids"
    _src_id_offset = 0
    _dst_id_offset = 0
    _edge_id_offset = 0

    for rank in range(partition_size):
        dst_nodes_in_each_partition[rank] = torch.nonzero(
            mapping_dst_ids_to_ranks == rank
        ).view(-1)
        src_nodes_in_each_partition[rank] = torch.nonzero(
            mapping_src_ids_to_ranks == rank
        ).view(-1)
        num_nodes = dst_nodes_in_each_partition[rank].numel()
        if num_nodes == 0:
            raise RuntimeError(
                f"Aborting partitioning, rank {rank} has 0 destination nodes to work on."
            )
        num_dst_nodes_in_each_partition[rank] = num_nodes

        num_nodes = src_nodes_in_each_partition[rank].numel()
        num_src_nodes_in_each_partition[rank] = num_nodes
        if num_nodes == 0:
            raise RuntimeError(
                f"Aborting partitioning, rank {rank} has 0 source nodes to work on."
            )

        # create mapping of global node IDs to/from "concatenated local" IDs
        ids = src_nodes_in_each_partition[rank]
        mapped_ids = torch.arange(
            start=_src_id_offset,
            end=_src_id_offset + ids.numel(),
            dtype=dtype,
            device=input_device,
        )
        _map_global_src_ids_to_local[ids] = mapped_ids - _src_id_offset
        graph_partition.map_global_src_ids_to_concatenated_local[ids] = mapped_ids
        graph_partition.map_concatenated_local_src_ids_to_global[mapped_ids] = ids
        _src_id_offset += ids.numel()

        ids = dst_nodes_in_each_partition[rank]
        mapped_ids = torch.arange(
            start=_dst_id_offset,
            end=_dst_id_offset + ids.numel(),
            dtype=dtype,
            device=input_device,
        )
        graph_partition.map_global_dst_ids_to_concatenated_local[ids] = mapped_ids
        graph_partition.map_concatenated_local_dst_ids_to_global[mapped_ids] = ids
        _dst_id_offset += ids.numel()

    graph_partition.num_src_nodes_in_each_partition = num_src_nodes_in_each_partition
    graph_partition.num_dst_nodes_in_each_partition = num_dst_nodes_in_each_partition

    # create local graph structures
    for rank in range(partition_size):
        offset_start = global_offsets[dst_nodes_in_each_partition[rank]].view(-1)
        offset_end = global_offsets[dst_nodes_in_each_partition[rank] + 1].view(-1)
        degree = offset_end - offset_start
        local_offsets = degree.view(-1).cumsum(dim=0)
        local_offsets = torch.cat(
            [
                torch.Tensor([0]).to(
                    dtype=dtype,
                    device=input_device,
                ),
                local_offsets,
            ]
        )

        partitioned_edge_ids = torch.cat(
            [
                torch.arange(
                    start=offset_start[i],
                    end=offset_end[i],
                    dtype=dtype,
                    device=input_device,
                )
                for i in range(len(offset_start))
            ]
        )

        ids = partitioned_edge_ids
        mapped_ids = torch.arange(
            _edge_id_offset,
            _edge_id_offset + ids.numel(),
            device=ids.device,
            dtype=ids.dtype,
        )
        graph_partition.map_global_edge_ids_to_concatenated_local[ids] = mapped_ids
        graph_partition.map_concatenated_local_edge_ids_to_global[mapped_ids] = ids
        _edge_id_offset += ids.numel()

        partitioned_src_ids = torch.cat(
            [
                global_indices[offset_start[i] : offset_end[i]].clone()
                for i in range(len(offset_start))
            ]
        )

        global_src_ids_on_rank, inverse_mapping = partitioned_src_ids.unique(
            sorted=True, return_inverse=True
        )
        remote_local_src_ids_on_rank = _map_global_src_ids_to_local[
            global_src_ids_on_rank
        ]

        _map_global_src_ids_to_local_graph = torch.zeros_like(mapping_src_ids_to_ranks)
        _num_local_indices = 0
        for rank_offset in range(partition_size):
            mask = mapping_src_ids_to_ranks[global_src_ids_on_rank] == rank_offset
            if partition_rank == rank_offset:
                # indices to send to this rank from this rank
                graph_partition.scatter_indices[rank] = (
                    remote_local_src_ids_on_rank[mask]
                    .detach()
                    .clone()
                    .to(dtype=torch.int64)
                )
            numel_mask = mask.sum().item()
            graph_partition.sizes[rank_offset][rank] = numel_mask

            tmp_ids = torch.arange(
                _num_local_indices,
                _num_local_indices + numel_mask,
                device=input_device,
                dtype=dtype,
            )
            _num_local_indices += numel_mask
            tmp_map = global_src_ids_on_rank[mask]
            _map_global_src_ids_to_local_graph[tmp_map] = tmp_ids

        local_indices = _map_global_src_ids_to_local_graph[partitioned_src_ids]
        graph_partition.num_indices_in_each_partition[rank] = local_indices.size(0)

        if rank == partition_rank:
            # local graph
            graph_partition.local_offsets = local_offsets
            graph_partition.local_indices = local_indices
            graph_partition.num_local_indices = graph_partition.local_indices.size(0)
            graph_partition.num_local_dst_nodes = num_dst_nodes_in_each_partition[rank]
            graph_partition.num_local_src_nodes = global_src_ids_on_rank.size(0)

            # partition-local mappings (local IDs to global)
            graph_partition.map_partitioned_src_ids_to_global = (
                src_nodes_in_each_partition[rank]
            )
            graph_partition.map_partitioned_dst_ids_to_global = (
                dst_nodes_in_each_partition[rank]
            )
            graph_partition.map_partitioned_edge_ids_to_global = partitioned_edge_ids

    for r in range(graph_partition.partition_size):
        err_msg = "error in graph partition: list containing sizes of exchanged indices does not match the tensor of indices to be exchanged"
        if (
            graph_partition.sizes[graph_partition.partition_rank][r]
            != graph_partition.scatter_indices[r].numel()
        ):
            raise AssertionError(err_msg)

    graph_partition = graph_partition.to(device=device)

    return graph_partition


def partition_graph_with_matrix_decomposition(
    global_offsets: torch.Tensor,
    global_indices: torch.Tensor,
    num_nodes: int,
    partition_book: torch.Tensor,
    partition_size: int,
    partition_rank: int,
    device: torch.device,
) -> GraphPartition:
    """
    Utility function which partitions a global graph given as CSC structure based on its adjacency
    matirx using 1-D row-wise decomposition. This approach ensures a 1D uniform distribution of nodes
    and their associated 1-hop incoming edges. By treating source and destination nodes equivalently
    during partitioning, this approach assumes the graph is not bipartite.
    This decomposition also ensures that the graph convolution (spMM) remains local by maintaining a copy of
    the local incoming edge features and the local node outputs from the graph convolution.
    The memory complexity of this approach is O[(N/P + E/P)*hid_dim*L], where N/E are the number of nodes/edges.
    The transformation from local node storage to local edge storage is achieved using nccl `alltoall`.

    Key differences from the existing graph partition scheme (partition_graph_with_id_mapping):
    (1) This function partitions the global node ID space uniformly, without distinguishing
    between source and destination nodes (i.e., matrix row ordering or column ordering). Both
    src/dst or row/col nodes are indexed consistently within the adjacency matrix.
    (2) Each local graph (sub-matrix) can be defined/constructed by just node/edge offsets from
    global graph.
    (3) The partitioning is performed on a global graph stored in CPU memory, and then each device
    (rank) constructs its local graph independently from the global csc matrix.

    Parameters
    ----------
    global_offsets : torch.Tensor
        CSC offsets, can live on the CPU
    global_indices : torch.Tensor
        CSC indices, can live on the CPU
    num_nodes : int
        number of nodes in the global graph
    partition_book : torch.Tensor
        the boundaries of 1-D row-decomp of adj. matrix for all ranks
    partition_size : int
        number of process groups across which graph is partitioned,
        i.e. the number of graph partitions
    partition_rank : int
        rank within process group managing the distributed graph, i.e.
        the rank determining which partition the corresponding local rank
        will manage
    device : torch.device
        device connected to the passed partition rank, i.e. the device
        on which the local graph and related buffers will live on
    """

    # initialize graph partition
    graph_partition = GraphPartition(
        partition_size=partition_size, partition_rank=partition_rank, device=device
    )
    dtype = global_indices.dtype
    # --------------------------------------------------------------
    # First partition the global row ptrs (dst nodes) to local row ptrs
    num_edges = global_indices.size(0)
    node_offset = partition_book[partition_rank]
    num_local_nodes = (
        partition_book[partition_rank + 1] - partition_book[partition_rank]
    )
    edge_partition_offset = global_offsets[node_offset]
    if node_offset + num_local_nodes > num_nodes:
        raise ValueError("Invalid node offset and number of local nodes")

    local_offsets = global_offsets[node_offset : node_offset + num_local_nodes + 1].to(
        device=device, non_blocking=True
    )
    graph_partition.local_offsets = local_offsets - edge_partition_offset
    graph_partition.num_local_dst_nodes = num_local_nodes

    # Scan through all partitions and compress the source nodes (edges) for each partition
    # to fill the local send/recv buffers for all-to-all communications
    partition_book = partition_book.to(device=device)
    for to_partition in range(partition_size):
        local_indices = global_indices[
            global_offsets[partition_book[to_partition]] : global_offsets[
                partition_book[to_partition + 1]
            ]
        ].to(device=device, non_blocking=True)
        # compress the columns (src nodes or local_indices) for each partition and record mapping (inverse_indices)
        global_src_node_at_partition, inverse_indices = local_indices.unique(
            sorted=True, return_inverse=True
        )
        global_src_node_at_partition_rank = (
            torch.bucketize(global_src_node_at_partition, partition_book, right=True)
            - 1
        )
        src_node_indices = torch.nonzero(
            global_src_node_at_partition_rank == partition_rank, as_tuple=False
        ).squeeze(1)
        # fill local send buffer for alltoalls (scatter selected nodes to_partition rank)
        graph_partition.scatter_indices[to_partition] = (
            global_src_node_at_partition[src_node_indices] - node_offset
        )
        # fill the numbers of indices (edges), dst nodes and src nodes for each partition
        graph_partition.num_indices_in_each_partition[to_partition] = (
            local_indices.size(0)
        )
        graph_partition.num_dst_nodes_in_each_partition[to_partition] = (
            partition_book[to_partition + 1] - partition_book[to_partition]
        )
        graph_partition.num_src_nodes_in_each_partition[to_partition] = (
            global_src_node_at_partition.size(0)
        )

        if to_partition == partition_rank:
            graph_partition.local_indices = inverse_indices
            graph_partition.num_local_indices = graph_partition.local_indices.size(0)
            graph_partition.num_local_src_nodes = global_src_node_at_partition.size(0)
            # map from local (compressed) column indices [0, ..., num_local_src_nodes] to their global node IDs
            graph_partition.map_partitioned_src_ids_to_global = (
                global_src_node_at_partition
            )

        for from_partition in range(partition_size):
            # fill all recv buffer sizes for alltoalls
            graph_partition.sizes[from_partition][to_partition] = torch.count_nonzero(
                global_src_node_at_partition_rank == from_partition
            )

    # trivial mappings due to 1D row-wise decomposition
    graph_partition.map_partitioned_dst_ids_to_global = torch.arange(
        node_offset, node_offset + num_local_nodes, dtype=dtype, device=device
    )
    graph_partition.map_partitioned_edge_ids_to_global = torch.arange(
        edge_partition_offset,
        edge_partition_offset + graph_partition.num_local_indices,
        dtype=dtype,
        device=device,
    )
    # trivial mappings due to 1D row-wise decomposition, with mem. cost O(E, N) at each dev; need to optimize
    graph_partition.map_concatenated_local_src_ids_to_global = torch.arange(
        num_nodes, dtype=dtype, device=device
    )
    graph_partition.map_concatenated_local_edge_ids_to_global = torch.arange(
        num_edges, dtype=dtype, device=device
    )
    graph_partition.map_concatenated_local_dst_ids_to_global = (
        graph_partition.map_concatenated_local_src_ids_to_global
    )
    graph_partition.map_global_src_ids_to_concatenated_local = (
        graph_partition.map_concatenated_local_src_ids_to_global
    )
    graph_partition.map_global_dst_ids_to_concatenated_local = (
        graph_partition.map_concatenated_local_src_ids_to_global
    )
    graph_partition.map_global_edge_ids_to_concatenated_local = (
        graph_partition.map_concatenated_local_edge_ids_to_global
    )
    graph_partition.matrix_decomp = True

    for r in range(graph_partition.partition_size):
        err_msg = "error in graph partition: list containing sizes of exchanged indices does not match the tensor of indices to be exchanged"
        if (
            graph_partition.sizes[graph_partition.partition_rank][r]
            != graph_partition.scatter_indices[r].numel()
        ):
            raise AssertionError(err_msg)
    graph_partition = graph_partition.to(device=device)
    return graph_partition


def partition_graph_nodewise(
    global_offsets: torch.Tensor,
    global_indices: torch.Tensor,
    partition_size: int,
    partition_rank: int,
    device: torch.device,
    matrix_decomp: bool = False,
) -> GraphPartition:
    """
    Utility function which partitions a global graph given as CSC structure naively
    by splitting both the IDs of source and destination nodes into chunks of equal
    size. For more details on partitioning in general see `GraphPartition`.
    The function performs the partitioning based on a global graph in CPU
    memory for each rank independently. It could be rewritten to e.g. only
    do it one rank and exchange the partitions or to an algorithm that also
    assumes an already distributed global graph, however, we expect global
    graphs to fit in CPU memory. After the partitioning, we can get rid off
    the larger one in CPU memory, only keep the local graphs on each GPU, and
    avoid tedious gather/scatter routines for exchanging partitions in the process.

    Parameters
    ----------
    global_offsets : torch.Tensor
        CSC offsets, can live on the CPU
    global_indices : torch.Tensor
        CSC indices, can live on the CPU
    partition_size : int
        number of process groups across which graph is partitioned,
        i.e. the number of graph partitions
    partition_rank : int
        rank within process group managing the distributed graph, i.e.
        the rank determining which partition the corresponding local rank
        will manage
    device : torch.device
        device connected to the passed partition rank, i.e. the device
        on which the local graph and related buffers will live on
    matrix_decomp : bool
        flag to enable matrix decomposition for partitioning
    """
    num_global_src_nodes = global_indices.max().item() + 1
    num_global_dst_nodes = global_offsets.size(0) - 1
    num_dst_nodes_per_partition = (
        num_global_dst_nodes + partition_size - 1
    ) // partition_size

    if matrix_decomp:
        if num_global_src_nodes != num_global_dst_nodes:
            raise ValueError(
                "Must be square adj. matrix (num_src=num_dst) for matrix decomposition"
            )
        partition_book = torch.arange(
            0,
            num_global_dst_nodes,
            num_dst_nodes_per_partition,
            dtype=global_indices.dtype,
        )
        partition_book = torch.cat(
            [
                partition_book,
                torch.tensor([num_global_dst_nodes], dtype=global_indices.dtype),
            ]
        )
        return partition_graph_with_matrix_decomposition(
            global_offsets,
            global_indices,
            num_global_dst_nodes,
            partition_book,
            partition_size,
            partition_rank,
            device,
        )

    num_src_nodes_per_partition = (
        num_global_src_nodes + partition_size - 1
    ) // partition_size

    mapping_dst_ids_to_ranks = (
        torch.arange(
            num_global_dst_nodes,
            dtype=global_offsets.dtype,
            device=global_offsets.device,
        )
        // num_dst_nodes_per_partition
    )
    mapping_src_ids_to_ranks = (
        torch.arange(
            num_global_src_nodes,
            dtype=global_offsets.dtype,
            device=global_offsets.device,
        )
        // num_src_nodes_per_partition
    )

    return partition_graph_with_id_mapping(
        global_offsets,
        global_indices,
        mapping_src_ids_to_ranks,
        mapping_dst_ids_to_ranks,
        partition_size,
        partition_rank,
        device,
    )


def partition_graph_by_coordinate_bbox(
    global_offsets: torch.Tensor,
    global_indices: torch.Tensor,
    src_coordinates: torch.Tensor,
    dst_coordinates: torch.Tensor,
    coordinate_separators_min: List[List[Optional[float]]],
    coordinate_separators_max: List[List[Optional[float]]],
    partition_size: int,
    partition_rank: int,
    device: torch.device,
) -> GraphPartition:
    """
    Utility function which partitions a global graph given as CSC structure.
    It partitions both the global ID spaces for source nodes and destination nodes
    based on their corresponding coordinates. Each partition will manage points which
    fulfill the boxconstraints specified by the specified coordinate separators. For each
    rank one is expected to specify the minimum and maximum coordinate value for each dimension.
    A partition the will manage all points for which ``min_val <= coord[d] < max_val`` holds. If one
    of the constraints is passed as `None`, it is assumed to be non-binding and the partition is defined
    by the corresponding half-space. Each rank maintains both a partition of the global source and
    destination nodes resulting from this subspace division.
    The function performs the partitioning based on a global graph in CPU
    memory for each rank independently. It could be rewritten to e.g. only
    do it one rank and exchange the partitions or to an algorithm that also
    assumes an already distributed global graph, however, we expect global
    graphs to fit in CPU memory. After the partitioning, we can get rid off
    the larger one in CPU memory, only keep the local graphs on each GPU, and
    avoid tedious gather/scatter routines for exchanging partitions in the process.
    Note: It is up to the user to ensure that the provided partition is valid.
    In particular, we expect each rank to receive a non-empty partition of node IDs.

    Examples
    --------
    >>> import torch
    >>> from physicsnemo.models.gnn_layers import partition_graph_by_coordinate_bbox
    >>> # simple graph with a degree of 2 per node
    >>> num_src_nodes = 8
    >>> num_dst_nodes = 4
    >>> offsets = torch.arange(num_dst_nodes + 1, dtype=torch.int64) * 2
    >>> indices = torch.arange(num_src_nodes, dtype=torch.int64)
    >>> # example with 2D coordinates
    >>> # assuming partitioning a 2D problem into the 4 quadrants
    >>> partition_size = 4
    >>> partition_rank = 0
    >>> coordinate_separators_min = [[0, 0], [None, 0], [None, None], [0, None]]
    >>> coordinate_separators_max = [[None, None], [0, None], [0, 0], [None, 0]]
    >>> device = "cuda:0"
    >>> # dummy coordinates
    >>> src_coordinates = torch.FloatTensor(
    ...     [
    ...         [-1.0, 1.0],
    ...         [1.0, 1.0],
    ...         [-1.0, -1.0],
    ...         [1.0, -1.0],
    ...         [-2.0, 2.0],
    ...         [2.0, 2.0],
    ...         [-2.0, -2.0],
    ...         [2.0, -2.0],
    ...     ]
    ... )
    >>> dst_coordinates = torch.FloatTensor(
    ...     [
    ...         [-1.0, 1.0],
    ...         [1.0, 1.0],
    ...         [-1.0, -1.0],
    ...         [1.0, -1.0],
    ...     ]
    ... )
    >>> # call partitioning routine
    >>> pg = partition_graph_by_coordinate_bbox(
    ...     offsets,
    ...     indices,
    ...     src_coordinates,
    ...     dst_coordinates,
    ...     coordinate_separators_min,
    ...     coordinate_separators_max,
    ...     partition_size,
    ...     partition_rank,
    ...     device,
    ... )
    >>> pg.local_offsets
    tensor([0, 2], device='cuda:0')
    >>> pg.local_indices
    tensor([0, 1], device='cuda:0')
    >>> pg.sizes
    [[0, 1, 1, 0], [0, 1, 1, 0], [1, 0, 0, 1], [1, 0, 0, 1]]
    >>>
    >>> # example with lat-long coordinates
    >>> # dummy coordinates
    >>> src_lat = torch.FloatTensor([-75, -60, -45, -30, 30, 45, 60, 75]).view(-1, 1)
    >>> dst_lat = torch.FloatTensor([-60, -30, 30, 30]).view(-1, 1)
    >>> src_long = torch.FloatTensor([-135, -135, 135, 135, -45, -45, 45, 45]).view(-1, 1)
    >>> dst_long = torch.FloatTensor([-135, 135, -45, 45]).view(-1, 1)
    >>> src_coordinates = torch.cat([src_lat, src_long], dim=1)
    >>> dst_coordinates = torch.cat([dst_lat, dst_long], dim=1)
    >>> # separate sphere at equator and 0 degree longitude into 4 parts
    >>> coordinate_separators_min = [
    ...     [-90, -180],
    ...     [-90, 0],
    ...     [0, -180],
    ...     [0, 0],
    ... ]
    >>> coordinate_separators_max = [
    ...     [0, 0],
    ...     [0, 180],
    ...     [90, 0],
    ...     [90, 180],
    ... ]
    >>> # call partitioning routine
    >>> partition_size = 4
    >>> partition_rank = 0
    >>> device = "cuda:0"
    >>> pg = partition_graph_by_coordinate_bbox(
    ...     offsets,
    ...     indices,
    ...     src_coordinates,
    ...     dst_coordinates,
    ...     coordinate_separators_min,
    ...     coordinate_separators_max,
    ...     partition_size,
    ...     partition_rank,
    ...     device,
    ... )
    >>> pg.local_offsets
    tensor([0, 2], device='cuda:0')
    >>> pg.local_indices
    tensor([0, 1], device='cuda:0')
    >>> pg.sizes
    [[2, 0, 0, 0], [0, 2, 0, 0], [0, 0, 2, 0], [0, 0, 0, 2]]

    Parameters
    ----------
    global_offsets : torch.Tensor
        CSC offsets, can live on the CPU
    global_indices : torch.Tensor
        CSC indices, can live on the CPU
    src_coordinates : torch.Tensor
        coordinates of each source node
    dst_coordinates : torch.Tensor
        coordinates of each destination node
    partition_size : int
        number of process groups across which graph is partitioned,
        i.e. the number of graph partitions
    partition_rank : int
        rank within process group managing the distributed graph, i.e.
        the rank determining which partition the corresponding local rank
        will manage
    device : torch.device
        device connected to the passed partition rank, i.e. the device
        on which the local graph and related buffers will live on
    """

    dim = src_coordinates.size(-1)
    if dst_coordinates.size(-1) != dim:
        raise ValueError()
    if len(coordinate_separators_min) != partition_size:
        a, b = len(coordinate_separators_min), partition_size
        error_msg = "Expected len(coordinate_separators_min) == partition_size"
        error_msg += f", but got {a} and {b} respectively"
        raise ValueError(error_msg)
    if len(coordinate_separators_max) != partition_size:
        a, b = len(coordinate_separators_max), partition_size
        error_msg = "Expected len(coordinate_separators_max) == partition_size"
        error_msg += f", but got {a} and {b} respectively"
        raise ValueError(error_msg)

    for rank in range(partition_size):
        if len(coordinate_separators_min[rank]) != dim:
            a, b = len(coordinate_separators_min[rank]), dim
            error_msg = f"Expected len(coordinate_separators_min[{rank}]) == dim"
            error_msg += f", but got {a} and {b} respectively"
        if len(coordinate_separators_max[rank]) != dim:
            a, b = len(coordinate_separators_max[rank]), dim
            error_msg = f"Expected len(coordinate_separators_max[{rank}]) == dim"
            error_msg += f", but got {a} and {b} respectively"

    num_global_src_nodes = global_indices.max().item() + 1
    num_global_dst_nodes = global_offsets.size(0) - 1

    mapping_dst_ids_to_ranks = torch.zeros(
        num_global_dst_nodes, dtype=global_offsets.dtype, device=global_offsets.device
    )
    mapping_src_ids_to_ranks = torch.zeros(
        num_global_src_nodes,
        dtype=global_offsets.dtype,
        device=global_offsets.device,
    )

    def _assign_ranks(mapping, coordinates):
        for p in range(partition_size):
            mask = torch.ones_like(mapping).to(dtype=torch.bool)
            for d in range(dim):
                min_val, max_val = (
                    coordinate_separators_min[p][d],
                    coordinate_separators_max[p][d],
                )
                if min_val is not None:
                    mask = mask & (coordinates[:, d] >= min_val)
                if max_val is not None:
                    mask = mask & (coordinates[:, d] < max_val)
            mapping[mask] = p

    _assign_ranks(mapping_src_ids_to_ranks, src_coordinates)
    _assign_ranks(mapping_dst_ids_to_ranks, dst_coordinates)

    return partition_graph_with_id_mapping(
        global_offsets,
        global_indices,
        mapping_src_ids_to_ranks,
        mapping_dst_ids_to_ranks,
        partition_size,
        partition_rank,
        device,
    )


class DistributedGraph:
    def __init__(
        self,
        global_offsets: torch.Tensor,
        global_indices: torch.Tensor,
        partition_size: int,
        graph_partition_group_name: str = None,
        graph_partition: Optional[GraphPartition] = None,
    ):  # pragma: no cover
        """
        Utility Class representing a distributed graph based on a given
        partitioning of a CSC graph structure. By default, a naive node-wise
        partitioning scheme is applied, see ``partition_graph_nodewise`` for
        details on that. This class then wraps necessary communication primitives
        to access all relevant feature buffers related to the graph.

        Parameters
        ----------
        global_offsets : torch.Tensor
            CSC offsets, can live on the CPU
        global_indices : torch.Tensor
            CSC indices, can live on the CPU
        partition_size : int
            Number of process groups across which graphs are distributed, expected to
            be larger than 1, i.e. an actual partition distributed among multiple ranks.
        partition_group_name : str, default=None
            Name of process group across which graphs are distributed. Passing no process
            group name leads to a parallelism across the default process group.
            Otherwise, the group size of a process group is expected to match partition_size.
        graph_partition : GraphPartition, optional
            Optional graph_partition, if passed as None, the naive
            node-wise partitioning scheme will be applied to global_offsets and global_indices,
            otherwise, these will be ignored and the passed partition will be used instead.
        """

        dist_manager = DistributedManager()
        self.device = dist_manager.device
        self.partition_rank = dist_manager.group_rank(name=graph_partition_group_name)
        self.partition_size = dist_manager.group_size(name=graph_partition_group_name)
        error_msg = f"Passed partition_size does not correspond to size of process_group, got {partition_size} and {self.partition_size} respectively."
        if self.partition_size != partition_size:
            raise AssertionError(error_msg)
        self.process_group = dist_manager.group(name=graph_partition_group_name)

        if graph_partition is None:
            # default partitioning scheme
            self.graph_partition = partition_graph_nodewise(
                global_offsets,
                global_indices,
                self.partition_size,
                self.partition_rank,
                self.device,
            )

        else:
            error_msg = f"Passed graph_partition.partition_size does not correspond to size of process_group, got {graph_partition.partition_size} and {self.partition_size} respectively."
            if graph_partition.partition_size != self.partition_size:
                raise AssertionError(error_msg)
            error_msg = f"Passed graph_partition.device does not correspond to device of this rank, got {graph_partition.device} and {self.device} respectively."
            if graph_partition.device != self.device:
                raise AssertionError(error_msg)
            self.graph_partition = graph_partition

        send_sizes = self.graph_partition.sizes[self.graph_partition.partition_rank]
        recv_sizes = [
            p[self.graph_partition.partition_rank] for p in self.graph_partition.sizes
        ]
        msg = f"GraphPartition(rank={self.graph_partition.partition_rank}, "
        msg += f"num_local_src_nodes={self.graph_partition.num_local_src_nodes}, "
        msg += f"num_local_dst_nodes={self.graph_partition.num_local_dst_nodes}, "
        msg += f"num_partitioned_src_nodes={self.graph_partition.num_src_nodes_in_each_partition[self.graph_partition.partition_rank]}, "
        msg += f"num_partitioned_dst_nodes={self.graph_partition.num_dst_nodes_in_each_partition[self.graph_partition.partition_rank]}, "
        msg += f"send_sizes={send_sizes}, recv_sizes={recv_sizes})"
        print(msg)

        dist.barrier(self.process_group)

    def get_src_node_features_in_partition(
        self,
        global_node_features: torch.Tensor,
        scatter_features: bool = False,
        src_rank: int = 0,
    ) -> torch.Tensor:  # pragma: no cover
        # if global features only on local rank 0 also scatter, split them
        # according to the partition and scatter them to other ranks

        if self.graph_partition.matrix_decomp:
            logger.warning(
                "Matrix decomposition assumes one type of node feature partition, and the graph"
                "adjacency matrix is square with identical src/dst node domains. "
                "So, only `get_dst_node_features_in_partition` is used/needed to get src or dst"
                "node features within a partition."
            )
            return self.get_dst_node_features_in_partition(
                global_node_features,
                scatter_features=scatter_features,
                src_rank=src_rank,
            )
        if scatter_features:
            global_node_features = global_node_features[
                self.graph_partition.map_concatenated_local_src_ids_to_global
            ]
            return scatter_v(
                global_node_features,
                self.graph_partition.num_src_nodes_in_each_partition,
                dim=0,
                src=src_rank,
                group=self.process_group,
            )

        return global_node_features.to(device=self.device)[
            self.graph_partition.map_partitioned_src_ids_to_global, :
        ]

    def get_src_node_features_in_local_graph(
        self, partitioned_src_node_features: torch.Tensor
    ) -> torch.Tensor:  # pragma: no cover
        # main primitive to gather all necessary src features
        # which are required for a csc-based message passing step
        return indexed_all_to_all_v(
            partitioned_src_node_features,
            indices=self.graph_partition.scatter_indices,
            sizes=self.graph_partition.sizes,
            use_fp32=True,
            dim=0,
            group=self.process_group,
        )

    def get_dst_node_features_in_partition(
        self,
        global_node_features: torch.Tensor,
        scatter_features: bool = False,
        src_rank: int = 0,
    ) -> torch.Tensor:  # pragma: no cover
        # if global features only on local rank 0 also scatter, split them
        # according to the partition and scatter them to other ranks
        if scatter_features:
            global_node_features = global_node_features.to(device=self.device)[
                self.graph_partition.map_concatenated_local_dst_ids_to_global
            ]
            return scatter_v(
                global_node_features,
                self.graph_partition.num_dst_nodes_in_each_partition,
                dim=0,
                src=src_rank,
                group=self.process_group,
            )

        return global_node_features.to(device=self.device)[
            self.graph_partition.map_partitioned_dst_ids_to_global, :
        ]

    def get_dst_node_features_in_local_graph(
        self,
        partitioned_dst_node_features: torch.Tensor,
    ) -> torch.Tensor:  # pragma: no cover
        # current partitioning scheme assumes that
        # local graph is built from partitioned IDs
        return partitioned_dst_node_features

    def get_edge_features_in_partition(
        self,
        global_edge_features: torch.Tensor,
        scatter_features: bool = False,
        src_rank: int = 0,
    ) -> torch.Tensor:  # pragma: no cover
        # if global features only on local rank 0 also scatter, split them
        # according to the partition and scatter them to other ranks
        if scatter_features:
            global_edge_features = global_edge_features[
                self.graph_partition.map_concatenated_local_edge_ids_to_global
            ]
            return scatter_v(
                global_edge_features,
                self.graph_partition.num_indices_in_each_partition,
                dim=0,
                src=src_rank,
                group=self.process_group,
            )

        return global_edge_features.to(device=self.device)[
            self.graph_partition.map_partitioned_edge_ids_to_global, :
        ]

    def get_edge_features_in_local_graph(
        self, partitioned_edge_features: torch.Tensor
    ) -> torch.Tensor:  # pragma: no cover
        # current partitioning scheme assumes that
        # local graph is built from partitioned IDs
        return partitioned_edge_features

    def get_global_src_node_features(
        self,
        partitioned_node_features: torch.Tensor,
        get_on_all_ranks: bool = True,
        dst_rank: int = 0,
    ) -> torch.Tensor:  # pragma: no cover
        error_msg = f"Passed partitioned_node_features.device does not correspond to device of this rank, got {partitioned_node_features.device} and {self.device} respectively."
        if partitioned_node_features.device != self.device:
            raise AssertionError(error_msg)

        if self.graph_partition.matrix_decomp:
            logger.warning(
                "Matrix decomposition assumes one type of node feature partition, and the graph"
                "adjacency matrix is square with identical src/dst node domains. "
                "So, only `get_global_dst_node_features` is used/needed to get global src or dst"
                "node features."
            )
            return self.get_global_dst_node_features(
                partitioned_node_features,
                get_on_all_ranks=get_on_all_ranks,
                dst_rank=dst_rank,
            )

        if not get_on_all_ranks:
            global_node_feat = gather_v(
                partitioned_node_features,
                self.graph_partition.num_src_nodes_in_each_partition,
                dim=0,
                dst=dst_rank,
                group=self.process_group,
            )
            if self.graph_partition.partition_rank == dst_rank:
                global_node_feat = global_node_feat[
                    self.graph_partition.map_global_src_ids_to_concatenated_local
                ]

            return global_node_feat

        global_node_feat = all_gather_v(
            partitioned_node_features,
            self.graph_partition.num_src_nodes_in_each_partition,
            dim=0,
            use_fp32=True,
            group=self.process_group,
        )
        global_node_feat = global_node_feat[
            self.graph_partition.map_global_src_ids_to_concatenated_local
        ]
        return global_node_feat

    def get_global_dst_node_features(
        self,
        partitioned_node_features: torch.Tensor,
        get_on_all_ranks: bool = True,
        dst_rank: int = 0,
    ) -> torch.Tensor:  # pragma: no cover
        error_msg = f"Passed partitioned_node_features.device does not correspond to device of this rank, got {partitioned_node_features.device} and {self.device} respectively."
        if partitioned_node_features.device != self.device:
            raise AssertionError(error_msg)

        if not get_on_all_ranks:
            global_node_feat = gather_v(
                partitioned_node_features,
                self.graph_partition.num_dst_nodes_in_each_partition,
                dim=0,
                dst=dst_rank,
                group=self.process_group,
            )
            if self.graph_partition.partition_rank == dst_rank:
                global_node_feat = global_node_feat[
                    self.graph_partition.map_global_dst_ids_to_concatenated_local
                ]

            return global_node_feat

        global_node_feat = all_gather_v(
            partitioned_node_features,
            self.graph_partition.num_dst_nodes_in_each_partition,
            dim=0,
            use_fp32=True,
            group=self.process_group,
        )
        global_node_feat = global_node_feat[
            self.graph_partition.map_global_dst_ids_to_concatenated_local
        ]
        return global_node_feat

    def get_global_edge_features(
        self,
        partitioned_edge_features: torch.Tensor,
        get_on_all_ranks: bool = True,
        dst_rank: int = 0,
    ) -> torch.Tensor:  # pragma: no cover
        error_msg = f"Passed partitioned_edge_features.device does not correspond to device of this rank, got {partitioned_edge_features.device} and {self.device} respectively."
        if partitioned_edge_features.device != self.device:
            raise AssertionError(error_msg)

        if not get_on_all_ranks:
            global_edge_feat = gather_v(
                partitioned_edge_features,
                self.graph_partition.num_indices_in_each_partition,
                dim=0,
                dst=dst_rank,
                group=self.process_group,
            )
            if self.graph_partition.partition_rank == dst_rank:
                global_edge_feat = global_edge_feat[
                    self.graph_partition.map_global_edge_ids_to_concatenated_local
                ]
            return global_edge_feat

        global_edge_feat = all_gather_v(
            partitioned_edge_features,
            self.graph_partition.num_indices_in_each_partition,
            dim=0,
            use_fp32=True,
            group=self.process_group,
        )
        global_edge_feat = global_edge_feat[
            self.graph_partition.map_global_edge_ids_to_concatenated_local
        ]
        return global_edge_feat