kahn.py

Go to the documentation of this file.

from itertools import dropwhile

from squander.gates.qgd_Circuit import qgd_Circuit as Circuit
from squander.partitioning.tools import build_dependency
import heapq

def kahn_partition(c, max_qubit, preparts=None):
    """
    Partition a circuit greedily using Kahn's algorithm

    Args:
        
        c: SQUANDER Circuit to partition
        
        max_qubit: Max qubits per partition
        
        preparts: Optional predefined partitions

    Returns:
        
        Partitioned circuit, parameter order (source_idx, dest_idx, param_count), partition assignments
    """
    # Build dependency graphs
    gate_dict, g, rg, gate_to_qubit, S = build_dependency(c)
    S = [(tuple(sorted(gate_to_qubit[x])), x) for x in S]
    heapq.heapify(S)
    
    L = []

    if preparts is None:
        def partition_condition(idx):
            return len(gate_to_qubit[S[idx][1]] | curr_partition) > max_qubit
    else:
        def partition_condition(idx):
            return not S[idx][1] in preparts[len(parts)-1]
    
    curr_partition = set()
    curr_idx = 0
    total = 0
    parts = [[]]

    while S:
        idx = next(dropwhile(partition_condition, range(len(S))), None)
        n = S[idx][1] if idx is not None else None
        if preparts is not None:
            assert (n is None) == (len(preparts[len(parts)-1]) == len(parts[-1])) #sanity check valid partitioning

        if n is None:  # partition cannot be expanded
            # Add partition to circuit
            total += len(parts[-1])
            # Reset for next partition
            curr_partition = set()
            parts.append([])
            idx = next(dropwhile(partition_condition, range(len(S))), None)
            n = S[idx][1] if idx is not None else None


        # Add gate to current partition
        parts[-1].append(n)
        curr_partition |= gate_to_qubit[n]
        curr_idx += gate_dict[n].get_Parameter_Num()

        # Update dependencies
        L.append(n)
        del S[idx]
        assert len(rg[n]) == 0
        for child in set(g[n]):
            g[n].remove(child)
            rg[child].remove(n)
            if not rg[child]:
                heapq.heappush(S, (tuple(sorted(gate_to_qubit[child])), child))

    # Add the last partition
    total += len(parts[-1])
    assert total == len(gate_dict)
    # print(parts)
    return kahn_partition_preparts(c, max_qubit, preparts=parts)


def kahn_partition_preparts(c, max_qubit, preparts):
    """
    Build partitioned circuit from predefined partitions

    Args:
        
        c: SQUANDER Circuit to partition
        
        max_qubit: Max qubits per partition
        
        preparts: Predefined partition assignments

    Returns:
        
        Partitioned circuit, parameter order (source_idx, dest_idx, param_count), partition assignments
    """
    top_circuit = Circuit(c.get_Qbit_Num())

    # Build dependency graphs
    gate_dict, g, rg, gate_to_qubit, S = build_dependency(c)
    S = [(tuple(sorted(gate_to_qubit[x])), x) for x in S]
    heapq.heapify(S)
    
    preparts = split_partition(preparts, g, rg)

    L = []
    param_order = []

    if preparts is None:
        def partition_condition(idx):
            return len(gate_to_qubit[S[idx][1]] | curr_partition) > max_qubit
    else:
        def partition_condition(idx):
            return not S[idx][1] in preparts[len(parts)-1]
    
    c = Circuit(c.get_Qbit_Num())
    curr_partition = set()
    curr_idx = 0
    total = 0
    parts = [[]]

    while S:
        idx = next(dropwhile(partition_condition, range(len(S))), None)
        n = S[idx][1] if idx is not None else None
        if preparts is not None:
            assert (n is None) == (len(preparts[len(parts)-1]) == len(parts[-1])) #sanity check valid partitioning

        if n is None:  # partition cannot be expanded
            # Add partition to circuit
            top_circuit.add_Circuit(c)
            total += len(c.get_Gates())
            # Reset for next partition
            curr_partition = set()
            c = Circuit(c.get_Qbit_Num())
            parts.append([])
            idx = next(dropwhile(partition_condition, range(len(S))), None)
            n = S[idx][1] if idx is not None else None


        # Add gate to current partition
        parts[-1].append(n)
        curr_partition |= gate_to_qubit[n]
        c.add_Gate(gate_dict[n])
        param_order.append((
            gate_dict[n].get_Parameter_Start_Index(), 
            curr_idx, 
            gate_dict[n].get_Parameter_Num()
        ))
        curr_idx += gate_dict[n].get_Parameter_Num()

        # Update dependencies
        L.append(n)
        del S[idx]
        assert len(rg[n]) == 0
        for child in set(g[n]):
            g[n].remove(child)
            rg[child].remove(n)
            if not rg[child]:
                heapq.heappush(S, (tuple(sorted(gate_to_qubit[child])), child))

    # Add the last partition
    top_circuit.add_Circuit(c)
    total += len(c.get_Gates())
    assert total == len(gate_dict)
    # print(parts)
    return top_circuit, param_order, parts


def split_partition(preparts, g, rg):
    """
    Split partitions into connected components

    Args:
        
        preparts: Initial partitions
        
        g: Forward dependency graph
        
        rg: Reverse dependency graph

    Returns:
        
        List of split partitions
    """
    L = []
    for part in preparts:
        V = set(part)
        while V:
            S = [next(iter(V))]
            L.append(set())
            while S:
                node = S.pop()
                if not node in V:
                    continue
                V.remove(node)
                L[-1].add(node)
                S.extend(rg[node])
                S.extend(g[node])
    ord_dict = {x: i for i, part in enumerate(preparts) for x in part}
    return [list(sorted(x, key=ord_dict.get)) for x in L]


def test_split_partition():
    """
    Test split_partition function for correctness.
    """
    q = 3
    c = Circuit(q)

    c.add_H(0)
    c.add_H(1)
    c.add_CNOT(0, 1)
    c.add_Z(2)
    c.add_H(2)

    expected = [[0, 1, 2], [3, 4]]

    _, g, rg, _, _ = build_dependency(c)

    split = split_partition([[0, 1, 2, 3, 4]], g, rg)
    
    print(split)

    assert { frozenset(x) for x in expected } == { frozenset(x) for x in split }


if __name__ == "__main__":
    test_split_partition()