import operator
from typing import Dict, List, Set, NamedTuple, Tuple

import torch
from torch.fx.passes.graph_manipulation import get_size_of_all_nodes
from torch.fx.experimental.partitioner_utils import (
    Partition,
    Device,
    PartitionerConfig,
    get_partition_to_latency_mapping,
    get_latency_of_partitioned_graph,
    NodeLatency,
    get_extra_size_of,
    PartitionMode,
)
from torch.fx.graph_module import GraphModule
from torch.fx.node import Node, map_arg
from torch.fx.passes.split_module import split_module


class DAGNode:
    """DAGNode class maintains useful information for a partition (submodule),
    and its input submodules and output submodules.
    """

    def __init__(
        self,
        submodule_node: Node,
        input_nodes: List[Node],
        output_nodes: List[Node],
        logical_device_ids: List[int],
        size_bytes: int,
    ) -> None:
        self.submodule_node: Node = submodule_node
        self.input_nodes: List[Node] = input_nodes
        self.output_nodes: List[Node] = output_nodes
        self.logical_device_ids: List[int] = logical_device_ids
        self.size_bytes = size_bytes

    def __str__(self) -> str:
        return str(self.submodule_node)


class DAG:
    """DAG class contains all the DAG nodes"""

    def __init__(self) -> None:
        self.nodes: List[DAGNode] = []

    def create_node(
        self,
        submodule_node: Node,
        input_nodes: List[Node],
        output_nodes: List[Node],
        logical_devices: List[int],
        size_bytes: int,
    ) -> None:
        node = DAGNode(
            submodule_node, input_nodes, output_nodes, logical_devices, size_bytes
        )
        self.nodes.append(node)


class PartitionResult(NamedTuple):
    """NameTuple used for returning DAG and a new fx module"""

    dag: DAG
    module_with_submodules: GraphModule


"""Followings are some helper functions for partition manipulation"""


def reset_partition_device(partitions):
    for partition in partitions:
        partition.logical_device_ids = []


def combine_two_partitions(
    partition_0: Partition, partition_1: Partition, partitions: List[Partition]
) -> None:
    """Given a list of partitions and its two partitions,
    combine these two partitions into a new one appending to the partitions
    and remove the previous two partitions from the list of partitions
    """
    partition = Partition(len(partitions))
    partition.nodes = partition_0.nodes.union(partition_1.nodes)
    partition.recalculate_mem_size()
    partitions.append(partition)
    partitions.remove(partition_0)
    partitions.remove(partition_1)
    reorganize_partitions(partitions)
    return


def set_parents_and_children(partitions: List[Partition]) -> None:
    """Given a list of partitions, mark parents and children for each partition"""
    # Go through all nodes in a partition.
    # If a node's user is in other partition,
    # then the other partition is this partition's children.
    # This partition is the other partition's parent
    for partition in partitions:
        partition.children = set()
        partition.parents = set()
    for partition in partitions:
        for node in partition.nodes:
            # For each node in the current partition, find its users
            users = node.users
            for n in users:
                # Find which the partition the user node belongs to.
                # Note that if the node itself is also belongs to that partition,
                # that partition is not the child of the current partition
                for p in partitions:
                    if p != partition and n in p.nodes and node not in p.nodes:
                        partition.children.add(p)
                        p.parents.add(partition)
    return


def reorganize_partitions(partitions: List[Partition]) -> None:
    """Given a list of partitions, reorganzie partiton id,
    its parents and its children for each partition
    """
    # Rearrange partition ids
    for i, partition in enumerate(partitions):
        partition.partition_id = i
    set_parents_and_children(partitions)
    return


def get_bfs_level_partition(partitions: List[Partition]) -> None:
    """Given a list of partitions,
    mark the bfs level for each partition
    """
    current_level: Set[Partition] = set()
    visited: Set[Partition] = set()
    for partition in partitions:
        # If a partition has no parent, it should be in root level
        if len(partition.parents) == 0:
            current_level.add(partition)
    next_level: Set[Partition] = set()
    level = 0
    # bfs
    while current_level:
        partition = current_level.pop()
        partition.bfs_level = level
        visited.add(partition)
        children = partition.children
        for child in children:
            if child not in next_level:
                next_level.add(child)
        if not current_level:
            current_level = next_level.copy()
            next_level = set()
            level += 1
    return


def get_node_to_partition_mapping(partitions: List[Partition]) -> Dict[Node, int]:
    """Given a list of partitions,return node to partition mapping"""
    node_to_partition: Dict[Node, int] = {}
    for partition in partitions:
        for node in partition.nodes:
            node_to_partition[node] = partition.partition_id
    return node_to_partition


def get_logical_id_to_device(devices: List[Device]) -> Dict[int, Device]:
    """Get a mapping from device logical ID to Device object."""
    logical_id_to_device: Dict[int, Device] = {}
    for d in devices:
        logical_id_to_device[d.logical_id] = d
    return logical_id_to_device


def get_device_partition_stats(
    partitions: List[Partition], devices: List[Device]
) -> Tuple[Dict[Device, List[Partition]], Dict[Device, int], List[Partition]]:
    """Given a list of partitions and a list of devices, returns:
    1. A mapping from device to partitions on it;
    2. A mapping from device to its remaining memory size;
    3. A list of partitions that do not have a device.
    """
    # logical id to device
    logical_id_to_device = get_logical_id_to_device(devices)
    # Track partitions on device
    device_to_partitions: Dict[Device, List[Partition]] = {}
    # Track device's left mem size
    device_to_left_mem_bytes: Dict[Device, int] = {}
    for d in devices:
        device_to_partitions[d] = []
        device_to_left_mem_bytes[d] = d.available_mem_bytes

    # Deal with the partitions that already have a device
    # and also collect all partitions without a device (no_device_partitions)
    no_device_partitions = []
    for partition in partitions:
        if partition.logical_device_ids != []:
            for logical_id in partition.logical_device_ids:
                device = logical_id_to_device[logical_id]
                device_to_partitions[device].append(partition)
                device_to_left_mem_bytes[device] -= partition.used_mem_bytes
        else:
            no_device_partitions.append(partition)

    return (
        device_to_partitions,
        device_to_left_mem_bytes,
        no_device_partitions,
    )


def get_device_to_partitions_mapping(
    partitions: List[Partition], devices: List[Device]
):
    """Given a list of partitions and a list of devices,
    map each partition into a device.
    """

    def calculate_extra_mem_bytes_needed_for(
        partition: Partition, partitions: List[Partition]
    ):
        all_nodes: Set[Node] = set()
        for p in partitions:
            all_nodes = all_nodes.union(p.nodes)
        if len(all_nodes) == 0:
            return partition.used_mem_bytes
        all_nodes = all_nodes.union(partition.nodes)
        extra_size_needed = 0
        for node in partition.nodes:
            extra_size_needed += get_extra_size_of(node, all_nodes)
        return extra_size_needed

    def find_device_for(partition: Partition):
        """Given a partition, find a logical device for the partition
        The algorithm is to put the partition on the device
        that has just enough mem left for that partition.
        device_to_left_mem_bytes is a dictionary between device and its left mem size
        sorted by its left mem size
        """
        for d in device_to_left_mem_bytes:
            extra_size_needed = calculate_extra_mem_bytes_needed_for(
                partition, device_to_partitions[d]
            )
            if extra_size_needed < device_to_left_mem_bytes[d]:
                device_to_partitions[d].append(partition)
                partition.logical_device_ids.append(d.logical_id)
                device_to_left_mem_bytes[d] -= extra_size_needed
                return True
        return False

    (
        device_to_partitions,
        device_to_left_mem_bytes,
        no_device_partitions,
    ) = get_device_partition_stats(partitions, devices)

    # Find devices for all the partitions without a device
    found_device = True
    for partition in no_device_partitions:
        device_to_left_mem_bytes = {
            d: left_mem_bytes
            for d, left_mem_bytes in sorted(
                device_to_left_mem_bytes.items(), key=lambda item: item[1]
            )
        }
        found_device = find_device_for(partition)
        if not found_device:
            break
    return found_device


def check_dependency(partition):
    """Given a partition,check if there is a circular dependency on
    this partition using bfs
    """
    visited: Set[Partition] = set([partition])
    queue: List[Partition] = [partition]
    while queue:
        p = queue.pop(0)
        for child in p.children:
            if child == partition:
                return True
            else:
                if child not in visited:
                    visited.add(child)
                    queue.append(child)
    return False


class Partitioner:
    """A fx module may not fit into one device.
    Partitioner class helps partition one fx module into submodules (partitions),
    so that the submodules can be executed crossing different accelerators.
    The main function of this class is self.partition_graph.
    It partitions the fx module based on the scheme specified in partition_config
    A DAG structure is returned
    along with a new fx module with submodule nodes.
    """

    def __init__(self) -> None:
        self.partitions: List[Partition] = []
        self.node_to_partition: Dict[Node, int] = {}
        self.devices: List[Device] = []

    def partition_graph(
        self,
        fx_module: GraphModule,
        torch_module: torch.nn.Module,
        partitioner_config: PartitionerConfig,
    ) -> PartitionResult:
        """Given the fx module, torch module and partitioner_config,
        find the partitions, do the partitions,
        and then return a DAG and a new fx module with submodule nodes (partitions)
        """
        self.graph_module = fx_module
        self.torch_module = torch_module
        self.devices = partitioner_config.devices
        if len(self.devices) == 0:
            raise RuntimeError("No devices")
        # Tag the size in bytes to all nodes in the graph_module.
        get_size_of_all_nodes(self.graph_module)
        # Check if there are op nodes in the fx module
        nodes = self.graph_module.graph.nodes
        if all(node.op in {"placeholder", "get_attr", "output"} for node in nodes):
            raise RuntimeError("No Partition since no operations in the module")
        # Calculate total size of the fx module
        total_size_of_graph = 0
        for node in nodes:
            if node.op == "output":
                break
            total_size_of_graph += node.size_bytes.total_size
        # Find the device with the max mem size
        device_with_max_mem = max(self.devices, key=lambda d: d.available_mem_bytes)
        # AOT based partition
        if partitioner_config.mode == PartitionMode.aot_based:
            self.aot_based_partition(
                partitioner_config.node_to_partition_mapping,
                partitioner_config.partition_to_logical_device_mapping,
            )
        # Single partition if the whole module can be fit into one device
        elif total_size_of_graph <= device_with_max_mem.available_mem_bytes:
            self.find_single_partition(
                total_size_of_graph, logical_device_id=device_with_max_mem.logical_id
            )
        elif total_size_of_graph > sum([d.available_mem_bytes for d in self.devices]):
            raise RuntimeError("Devices have no enough memory for the module")
        else:
            # Sparse nn based partition
            if partitioner_config.mode == PartitionMode.sparse_nn:
                available_mem_bytes = self.devices[0].available_mem_bytes
                if not all(
                    device.available_mem_bytes == available_mem_bytes
                    for device in self.devices
                ):
                    raise RuntimeError("All devices must have same memory size!")
                # sparse_nn_partition only support same memory size
                # TODO: add different size support for sparse_nn_partition
                self.sparse_nn_partition(available_mem_bytes)
            # Cost aware partition
            elif partitioner_config.mode == PartitionMode.cost_aware:
                self.cost_aware_partition(
                    partitioner_config.transfer_rate_bytes_per_sec,
                    partitioner_config.node_to_latency_mapping,
                )
            # KL based partition
            elif partitioner_config.mode == PartitionMode.kl_based:
                self.kl_based_partition(
                    partitioner_config.transfer_rate_bytes_per_sec,
                    partitioner_config.node_to_latency_mapping,
                )
            else:
                self.size_based_partition()

        # Saturate host if possible.
        if partitioner_config.saturate_host:
            self.saturate_host()

        # Partition the graph module based on the partition assignment.
        module_with_submodules = self.do_partition()

        # The DAG contains DAGNodes with info of each partition's input nodes, output nodes
        # and how partitions are connected.
        dag = self.dump_dag(module_with_submodules)
        ret = PartitionResult(dag, module_with_submodules)
        return ret

    def find_single_partition(
        self, total_size_of_graph, logical_device_id: int = 0
    ) -> None:
        """Fit the whole fx module into one device"""
        partition_0 = self.create_partition()
        for node in self.graph_module.graph.nodes:
            if node.op == "output":
                # Skip the output node, but there can
                # be nodes after the output in certain cases.
                continue
            partition_0.nodes.add(node)
        partition_0.used_mem_bytes = total_size_of_graph
        partition_0.logical_device_ids = [logical_device_id]
        # Get the node to partition mapping
        self.node_to_partition = get_node_to_partition_mapping(self.partitions)
        return

    def size_based_partition(self) -> None:
        """This method is to partition the fx module based on memory size.
        It uses greedy approach. The result may not be the best.
        The basic idea is:
        Step 1:
        Find a device which has enough memory to fit the current node, create a empty partition
        with the size of that device.
        Then keep adding the following nodes into the partition until the partition is full.
        Step 2:
        Repeat Step 1 until no device left
        Step 3:
        If some nodes are left, create a partition for each left node (single node partition).
        and then try to map those partitions into logical devices with enough mem left.
        """

        def find_device_based_on_size(node) -> Device:
            """Given a node, this function is to find a logical device
            that could fit the node.
            """
            mem_size_needed = get_extra_size_of(node, set())
            device = Device("", -1, -1)
            for d in self.devices:
                if (
                    d not in occupied_devices
                    and d.available_mem_bytes >= mem_size_needed
                ):
                    device = d
                    break
            if device.available_mem_bytes < 0:
                raise RuntimeError(str(node) + "is too large to fit any device")
            occupied_devices.append(device)
            return device

        # Track partition and its left mem size
        partition_to_left_mem_bytes: Dict[Partition, int] = {}
        # Track all the devices that have been used
        occupied_devices: List[Device] = []
        partition = self.create_partition()
        for node in self.graph_module.graph.nodes:
            if node.op in {"call_module", "call_method", "call_function"}:
                # Check if there are devices left
                if len(self.partitions) <= len(self.devices):
                    total_size_of_input_nodes = get_extra_size_of(node, partition.nodes)
                    # Check if the current partition is the very first partition
                    if partition.used_mem_bytes == 0:
                        # Find a device to fit the first node, return available mem size
                        device = find_device_based_on_size(node)
                        occupied_devices.append(device)
                        # Update partition and its left mem size
                        partition_to_left_mem_bytes[
                            partition
                        ] = device.available_mem_bytes
                        # Update available mem for the current partition
                        partition.logical_device_ids.append(device.logical_id)
                    else:
                        # The current partition is not the first partition
                        # Check if the current node can fit into current partition
                        if (
                            partition_to_left_mem_bytes[partition]
                            < total_size_of_input_nodes
                        ):
                            # Check if no device is left
                            if len(self.partitions) == len(self.devices):
                                # No device is left
                                # Put the previous partitions into a list (non_single_node_partitions)
                                non_single_node_partitions = self.partitions[:]
                                # Create the first single node partition for the current node
                                self.create_single_node_partition(node)
                                continue
                            # Some devices are still left
                            # Create a new partition with a mem size that is enough for the current node
                            device = find_device_based_on_size(node)
                            partition = self.create_partition()
                            total_size_of_input_nodes = get_extra_size_of(
                                node, partition.nodes
                            )
                            partition_to_left_mem_bytes[
                                partition
                            ] = device.available_mem_bytes
                            partition.logical_device_ids.append(device.logical_id)
                    partition.add_node(node)
                    partition_to_left_mem_bytes[partition] -= total_size_of_input_nodes
                # Create single node partitions if no device is left
                else:
                    self.create_single_node_partition(node)
        reorganize_partitions(self.partitions)
        # Get the node to partition mapping
        self.node_to_partition = get_node_to_partition_mapping(self.partitions)
        # Mapping all partitions into device
        found_partition_to_device_mapping = get_device_to_partitions_mapping(
            self.partitions, self.devices
        )
        if not found_partition_to_device_mapping:
            raise RuntimeError("Cannot Get a Valid Partition to Logical Device Mapping")
        return

    def saturate_host(self) -> None:
        """Saturate host by assigning replicates to unused devices with enough memory.
        It uses a greedy approach to find a next available set of devices to place all split
        partitions: For each used device, it searches for an idle device with minimal memory
        size that can hold all the partition located on that device; If the search is successful
        for all used devices, it then assigns the new devices' logical ID to the corresponding
        partition.
        """
        (
            device_to_partitions,
            device_to_left_mem_bytes,
            no_device_partitions,
        ) = get_device_partition_stats(self.partitions, self.devices)

        assert (
            len(no_device_partitions) == 0
        ), f"Expect no_device_partitions has 0 device, but get {len(no_device_partitions)}"

        # Devices that hold partitions
        used_devices = [d for d in self.devices if len(device_to_partitions[d]) > 0]
        # Track replicates of the assigned devices
        replicated_device_to_used_device: Dict[Device, Device] = {}

        while len(used_devices) * 2 + len(replicated_device_to_used_device) <= len(
            self.devices
        ):
            # Success flag for this round
            success = True
            # Devices that have not been assigned
            idle_devices = [
                d
                for d in self.devices
                if d not in used_devices and d not in replicated_device_to_used_device
            ]
            # Temporary mapping from replicated device to original device
            temp_replicate_mapping = {}

            # Find a new device to replicate all partitions on an used device
            for used_device in used_devices:
                # Idle devices that have enough memory
                available_devices = [
                    d
                    for d in idle_devices
                    if d.available_mem_bytes
                    >= used_device.available_mem_bytes
                    - device_to_left_mem_bytes[used_device]
                ]
                if len(available_devices) == 0:
                    success = False
                    break
                new_device = min(available_devices, key=lambda d: d.available_mem_bytes)
                idle_devices.remove(new_device)
                temp_replicate_mapping[new_device] = used_device

            if not success:
                break
            replicated_device_to_used_device.update(temp_replicate_mapping)

        # Update logical device IDs assigned to the partitions
        for (
            replicate_device,
            original_device,
        ) in replicated_device_to_used_device.items():
            logical_id = replicate_device.logical_id
            for partition in device_to_partitions[original_device]:
                partition.logical_device_ids.append(logical_id)
        for p in self.partitions:
            print(p.logical_device_ids)

    def do_partition(self) -> GraphModule:
        """Return a new fx module with submodule nodes (partitions)."""
        module_with_submodules = split_module(
            self.graph_module,
            self.torch_module,
            lambda node: self.node_to_partition[node],
        )
        return module_with_submodules

    def dump_dag(self, module_with_submodules: GraphModule) -> DAG:
        """Return the dag structure and the new fx module with submodules."""
        dag = DAG()
        for node in module_with_submodules.graph.nodes:
            if node.op == "output":
                break
            if node.op in {"placeholder", "get_attr"}:
                continue
            if node.target == operator.__getitem__:
                continue
            input_nodes: Dict[Node, None] = {}
            map_arg(node.args, lambda n: input_nodes.setdefault(n))
            map_arg(node.kwargs, lambda n: input_nodes.setdefault(n))
            # When a node has two or more output nodes,
            # it outputs its result to 'getitem' nodes.
            # Those 'getitem' nodes are the output node for this node.
            # Otherwise, the output node is this node itself.
            if len(node.users) > 1:
                output_nodes = list(node.users)
            else:
                output_nodes = [node]
            partition_id = int(node.name.rsplit("_", 1)[-1])
            device_ids = self.partitions[partition_id].logical_device_ids
            size_bytes = self.partitions[partition_id].used_mem_bytes
            dag.create_node(
                node, list(input_nodes), output_nodes, device_ids, size_bytes
            )
        return dag

    def create_partition(self) -> Partition:
        """Create a partition and append it to self.partitions."""
        partition_id = len(self.partitions)
        partition = Partition(partition_id)
        self.partitions.append(partition)
        return partition

    def create_single_node_partition(self, node):
        """Create a partition for a single node"""
        partition = self.create_partition()
        partition.add_node(node)
        return

    def sparse_nn_partition(self, available_mem_bytes: int) -> None:
        """This method partition a sparse nn module.
        It is size based partition but different from size_based_partition,
        it only works when all the devices have same memory size (available_mem_bytes).
        In the future, devices with different mem sizes will be supported like size_based_partition.
        It first traverse all the nodes and do the partitions based on the same memory size.
        If the current partition has no enough memory left for a new op node
        (call_module, call_method, call_function), a new partition is created.
        When crossing the boundary between non-embedding nodes and embedding nodes,
        a new partition is created regardlessly.
        For example, if the current node is a non-embedding node but the next node is an
        embedding node, a new partition is created for the next node.
        After the partition, the partitions are combined as much as possible.
        The rule is that a non-embedding partition only
        combines with another non-embedding one.
        So as the embedding partitions.
        """

        def combine_partitions_based_on_size(
            partitions: List[Partition], available_mem_bytes: int
        ) -> None:
            """Combining small partitions together to keep as less partitions as possible.
            Here is an example of the algorithm to do this:
            Assume some partitions, we first sort them based on partiiton used memory size.
            [(partition_4, 1), (partition_3, 1), (partition_2, 2), (partition_1, 7), (partition_0, 9)]
            The available memory is 10.
            step 1: self.find_partition_to_combine_based_on_size()
            First, mark bfs level for each partition
            Second, look the smallest partition, partition_4: 10 - 1 = 9
            It means any partition has a used memory equal or less than 9 could combine this partition
            We go from the largest and selection partition_0.
            Check the bfs level for two partitions, if the level difference is less than 2,
            it can be combined.
            step 2: repeat step 1 until no partitions can be combined
            """
            find_combination = True
            while find_combination:
                # Sort partitions based on memory size
                sorted_partitions = sorted(partitions, key=lambda p: p.used_mem_bytes)
                # Mark bfs level
                get_bfs_level_partition(self.partitions)
                find_combination, partitions = find_partition_to_combine_based_on_size(
                    sorted_partitions, available_mem_bytes, partitions
                )
            return

        def calculate_mem_bytes_needed(p1, p2):
            """Given two partitions, calculate how many mem bytes
            are needed if two partitions are combined
            """
            nodes = p1.nodes.union(p2.nodes)
            mem_bytes_needed = 0
            for node in nodes:
                mem_bytes_needed += get_extra_size_of(node, nodes)
            return mem_bytes_needed

        def find_partition_to_combine_based_on_size(
            sorted_partitions: List[Partition],
            available_mem_bytes: int,
            partitions: List[Partition],
        ) -> Tuple[bool, List[Partition]]:
            """step 1 in combine_partition_based_on_size()"""
            find_combination = False
            smallest_partition = sorted_partitions.pop(0)
            for p in sorted_partitions[::-1]:
                if abs(smallest_partition.bfs_level - p.bfs_level) <= 1:
                    # Calculate how many bytes needed if combined
                    mem_bytes_needed = calculate_mem_bytes_needed(p, smallest_partition)
                    if mem_bytes_needed <= available_mem_bytes:
                        combine_two_partitions(p, smallest_partition, self.partitions)
                        partitions.remove(smallest_partition)
                        partitions.remove(p)
                        partitions.append(self.partitions[-1])
                        find_combination = True
                        break
            return find_combination, partitions

        def reset_partition_in_sparse_nn(partition, new_partition=True):
            """If crossing the boudary between non-embedding nodes and
            embedding nodes, create a new partition
            """
            if in_embedding_region:
                embedding_partitions.append(partition)
            else:
                non_embedding_partitions.append(partition)
            if new_partition:
                partition = self.create_partition()
                partition.left_mem_bytes = available_mem_bytes
                return partition
            return None

        def is_embedding_node(node: Node) -> bool:
            """Check if a node is an embedding node"""
            if node.op == "call_module":
                submodule = self.graph_module
                for atom in str(node.target).split("."):
                    if not hasattr(submodule, atom):
                        raise RuntimeError(
                            f"Module {submodule} has no attribute {atom}"
                        )
                    submodule = getattr(submodule, atom)
                    if "Embedding" in str(submodule):
                        return True
            return False

        # Track embedding partitons and non-embedding partitions separately
        embedding_partitions: List[Partition] = []
        non_embedding_partitions: List[Partition] = []
        # A Flag to check the boundary
        in_embedding_region: bool = False
        partition = self.create_partition()
        for node in self.graph_module.graph.nodes:
            if node.op in {"call_module", "call_method", "call_function"}:
                # Check if crossing the boundary between embedding nodes and non embedding nodes
                if is_embedding_node(node) != in_embedding_region:
                    # Crossing the boundary
                    # Check if the current partition is an empty partition
                    if partition.used_mem_bytes != 0:
                        # The current partition isn't an empty partition. Create a new one.
                        partition = reset_partition_in_sparse_nn(partition)
                    in_embedding_region = not in_embedding_region
                total_size_of_input_nodes = get_extra_size_of(node, partition.nodes)
                if (
                    total_size_of_input_nodes + partition.used_mem_bytes
                    > available_mem_bytes
                ):
                    partition = reset_partition_in_sparse_nn(partition)
                    total_size_of_input_nodes = get_extra_size_of(node, partition.nodes)
                    if total_size_of_input_nodes > available_mem_bytes:
                        raise RuntimeError(
                            node.target + "is too large to fit into a device"
                        )
                partition.add_node(node)
        reset_partition_in_sparse_nn(partition, new_partition=False)
        # Set parents and children for partitions
        set_parents_and_children(self.partitions)
        # Combining non-embedding partitions
        combine_partitions_based_on_size(non_embedding_partitions, available_mem_bytes)
        # Combining embedding partitions
        combine_partitions_based_on_size(embedding_partitions, available_mem_bytes)
        total_size_of_non_embedding_partitions = 0
        for partition in non_embedding_partitions:
            total_size_of_non_embedding_partitions += partition.used_mem_bytes
        # Check if devices are enough for all partitions
        if len(embedding_partitions) > len(self.devices):
            msg = (
                "Need "
                + str(len(embedding_partitions))
                + " devices, but only "
                + str(len(self.devices))
                + " provided"
            )
            raise RuntimeError(msg)
        occupied_devices = []
        for i, partition in enumerate(embedding_partitions):
            # Check if all non-embedding partitions can fit into embedding partition devices
            if (
                total_size_of_non_embedding_partitions + partition.used_mem_bytes
                > available_mem_bytes
            ):
                raise RuntimeError(
                    "partition_"
                    + str(partition.partition_id)
                    + "(embedding partition) and non embedding partitions can not fit into one device"
                )
            else:
                # Add logical device to the partition
                partition.logical_device_ids = [self.devices[i].logical_id]
                occupied_devices.append(self.devices[i].logical_id)
        # Add logical devices to the non_embedding_partitions
        for partition in non_embedding_partitions:
            partition.logical_device_ids = occupied_devices
        # Get the node to partition mapping
        self.node_to_partition = get_node_to_partition_mapping(self.partitions)
        return

    def cost_aware_partition(
        self,
        transfer_rate_bytes_per_sec: float,
        node_to_latency_mapping: Dict[Node, NodeLatency],
    ) -> None:
        """This method is to partition the fx module based on the cost.
        The cost is the total latency of running the whole fx module.
        In partitioner_utils.py, the cost model is built.
        The cost aware partition algorithm is:
        #1. At every begining, each node is a partition.
            Then we map all the partitions to the devices
            and calculate the cost
        #2. Then try to pre-combine any two of the partitions if the two
            partitions can be combined.
            (the bfs level is less than 2 or two partitions are connected and
            can find partition to device mapping)
            See if any partition pair could reduce the current cost.
            Choose the pair that shows the minimum cost and then combine them
        #3. Repeat #2 until the cost cannot be reduced.
        """

        def try_combining_partitions(p0_index, p1_index, partitions) -> float:
            """Given two partitions and a list of partitions, combine these two partitions
            and see what is the cost of the modified partition list
            """
            p0 = partitions[p0_index]
            p1 = partitions[p1_index]
            """If two partitions' bfs level are less than 2 or two partitions are connected to each other,
               then they can be combined
            """
            if (
                (abs(p0.bfs_level - p1.bfs_level) <= 1)
                or (p0 in p1.parents)
                or p0 in (p1.children)
            ):
                combine_two_partitions(p0, p1, partitions)
                # Check if a circular dependency exists after combining
                if check_dependency(partitions[-1]):
                    return float("inf")
                # Check if the modified partition list can be mapped to devices after combination
                reset_partition_device(partitions)
                found_deivce = get_device_to_partitions_mapping(
                    partitions, self.devices
                )
                if not found_deivce:
                    return float("inf")
                # Calculate the new cost
                partition_to_latency_mapping = get_partition_to_latency_mapping(
                    partitions, node_to_latency_mapping
                )
                cost = get_latency_of_partitioned_graph(
                    partitions,
                    partition_to_latency_mapping,
                    transfer_rate_bytes_per_sec,
                )
                return cost
            # If two partition can not be combined, the cost is inf
            return float("inf")

        def search_combination(
            transfer_rate_bytes_per_sec, node_to_latency_mapping
        ) -> bool:
            """Given transfer rate between partitions and each node's latency,
            find two partitions to combine so the cost of the partitions can
            be reduced.
            The algorithm is :
            1. Go through all the partition pairs and see
            if any pair of partitions can be combined.
            2. Calculate the cost after the combination.
            3. Select the minimum cost and combine its cooresponding partition pair.
            """
            partition_to_latency_mapping = get_partition_to_latency_mapping(
                self.partitions, node_to_latency_mapping
            )
            cost = get_latency_of_partitioned_graph(
                self.partitions,
                partition_to_latency_mapping,
                transfer_rate_bytes_per_sec,
            )
            if len(self.partitions) == 1:
                return False
            partition_pair: List[int] = []
            for i in range(len(self.partitions) - 1):
                for j in range(i + 1, len(self.partitions)):
                    # Try to combine the partition pair
                    # and see the new cost after combination
                    new_cost = try_combining_partitions(i, j, self.partitions[:])
                    if new_cost <= cost:
                        partition_pair = [i, j]
                        cost = new_cost
                    reorganize_partitions(self.partitions)
            # If a partition pair is found, combine them
            if len(partition_pair) != 0:
                p0 = self.partitions[partition_pair[0]]
                p1 = self.partitions[partition_pair[1]]
                combine_two_partitions(p0, p1, self.partitions)
            get_bfs_level_partition(self.partitions)
            reset_partition_device(self.partitions)
            get_device_to_partitions_mapping(self.partitions, self.devices)
            return len(partition_pair) != 0

        for node in self.graph_module.graph.nodes:
            if node.op not in {"placeholder", "get_attr", "output"}:
                self.create_single_node_partition(node)
        # Set up parent partitions and children partitions for each partition
        set_parents_and_children(self.partitions)
        # Get bfs level for each partition
        get_bfs_level_partition(self.partitions)
        find_combination = True
        while find_combination:
            # Search for a pair partition to generate the minimum new cost,
            # then combine them
            find_combination = search_combination(
                transfer_rate_bytes_per_sec, node_to_latency_mapping
            )
        # Make sure all partitions are set up correctly
        reorganize_partitions(self.partitions)
        # Set up node to partition mapping
        self.node_to_partition = get_node_to_partition_mapping(self.partitions)
        return

    def kl_based_partition(
        self,
        transfer_rate_bytes_per_sec: float,
        node_to_latency_mapping: Dict[Node, NodeLatency],
    ) -> None:
        """This function is a cost aware partition based
        on Kernighan-Lin algorithm.
        First, the graph is partitioned using size_based_partition.
        Then, each node is swapped with any other node in a different
        partition, and at the same time, the cost is estimated after
        the swapping.
        For example, we have nodes n0, n1, n2, n3 and n4.
        Using size_based_partition, n0 and n1 are in Partition p0.
        n2, n3 and n4 in Partition p1. The current cost is esimated.
        We first tried using n0 to swap with n2 from the other partiton.
        Then we see that swapping n0 and n2 shows a lower cost
        than the current cost and it is the minimum among other pairs like
        (n0, None)(This means moving n0 to Partition without swapping other nodes),
        (n0, n3) and (n0, n4). We swap n0 and n2 and set the new cost
        as the current cost.
        Then We repeat this process for all the other nodes until all swapping pairs
        are tried.
        """

        def swap_nodes(n0, n1, p0, p1):
            # Either n0 or n1 could be None
            # That means we simply move the node
            # to another partition
            if n0 is not None:
                p0.remove_node(n0)
                p1.add_node(n0)
            if n1 is not None:
                p0.add_node(n1)
                p1.remove_node(n1)

        def try_swap_nodes(
            n0, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_sec
        ):
            cost = float("inf")
            swap_nodes(n0, n1, p0, p1)
            # Reorganize partitions after swapping
            reorganize_partitions(self.partitions)
            # Check if there is a circular dependency after swapping
            if (not check_dependency(p0)) and (not check_dependency(p1)):
                reset_partition_device(self.partitions)
                partition_to_latency_mapping = get_partition_to_latency_mapping(
                    self.partitions, node_to_latency_mapping
                )
                # Check if all partitions can be mapped to logical devices after swapping
                found_device = get_device_to_partitions_mapping(
                    self.partitions, self.devices
                )
                if not found_device:
                    cost = float("inf")
                else:
                    cost = get_latency_of_partitioned_graph(
                        self.partitions,
                        partition_to_latency_mapping,
                        transfer_rate_bytes_per_sec,
                    )
            # Swap back and reset all partitions back to original
            swap_nodes(n1, n0, p0, p1)
            reorganize_partitions(self.partitions)
            reset_partition_device(self.partitions)
            get_device_to_partitions_mapping(self.partitions, self.devices)
            return cost

        def swap_node_to_partition(
            node, p0, p1, node_to_latency_mapping, transfer_rate_per_sec
        ):
            """This function helps to swap one node from partition p0
            with all the nodes in another partition p1
            """
            p1_nodes = list(p1.nodes) + [None]
            min_cost = float("inf")
            node_pair: List[Node] = []
            for n1 in p1_nodes:
                # Ignore the node if it is not a op node
                if n1 is not None and n1.op in {"placeholder", "get_attr"}:
                    continue
                # Try swapping node in p0 with n1 in p1
                cost = try_swap_nodes(
                    node, n1, p0, p1, node_to_latency_mapping, transfer_rate_per_sec
                )
                if cost < min_cost:
                    node_pair = [node, n1]
                    min_cost = cost
            return cost, node_pair

        # First use size_base_partition
        self.size_based_partition()
        partition_to_latency_mapping = get_partition_to_latency_mapping(
            self.partitions, node_to_latency_mapping
        )
        # Calculate the cost of the partitions
        cost = get_latency_of_partitioned_graph(
            self.partitions, partition_to_latency_mapping, transfer_rate_bytes_per_sec
        )
        # Keep tracking the node pair that shows the better cost
        node_pair: List[Node] = []
        # Keep tracking the partition pair of node pair
        partition_pair: List[Partition] = []
        # Collect all the op nodes from the graph
        op_nodes = []
        for n in self.graph_module.graph.nodes:
            if n.op not in {"placeholder", "get_attr", "output"}:
                op_nodes.append(n)
        for node in op_nodes:
            # Find which partition the current node belongs
            p0_index = self.node_to_partition[node]
            p0 = self.partitions[p0_index]
            # Go through all the other partitions to swap
            # with other nodes from those partitions
            for p1_index, _ in enumerate(self.partitions):
                if p0_index != p1_index:
                    p1 = self.partitions[p1_index]
                    new_cost, new_node_pair = swap_node_to_partition(
                        node,
                        p0,
                        p1,
                        node_to_latency_mapping,
                        transfer_rate_bytes_per_sec,
                    )
                    # Update the cost
                    # Track the swapped node pair and their partitions
                    if new_cost < cost:
                        cost = new_cost
                        node_pair = new_node_pair
                        partition_pair = [p0, p1]
            # Do the swapping after trying all the nodes from a partition
            if len(node_pair) != 0:
                swap_nodes(
                    node_pair[0], node_pair[1], partition_pair[0], partition_pair[1]
                )
                reorganize_partitions(self.partitions)
                get_device_to_partitions_mapping(self.partitions, self.devices)
        reorganize_partitions(self.partitions)
        # Mapping the device to the partition
        get_device_to_partitions_mapping(self.partitions, self.devices)
        return

    def aot_based_partition(
        self, node_to_partition_mapping, partition_to_logical_device_mapping
    ):
        """This function helps to rebuild the partitions given the nodes and its
        corresponding partition id
        """
        partition_id_to_partition_mapping: Dict[int, Partition] = {}
        self.node_to_partition = node_to_partition_mapping
        for node in self.node_to_partition:
            partition_id = self.node_to_partition[node]
            # If the requested partition has not been created, create the partition
            if partition_id not in partition_id_to_partition_mapping:
                partition = Partition(partition_id)
                self.partitions.append(partition)
                partition_id_to_partition_mapping[partition_id] = partition
                partition.logical_device_ids = partition_to_logical_device_mapping[
                    partition_id
                ]
            else:
                partition = partition_id_to_partition_mapping[
                    self.node_to_partition[node]
                ]
            # Add the current node into the partition
            partition.add_node(node)