make_diagram.py

import cairocffi as cairo
import collections
from itertools import chain
from nengo_spinnaker.partition import divide_slice
from nengo_spinnaker.utils.itertools import flatten
import numpy as np
import pygraphviz as pgv
import random
from rig.machine import Cores, Machine
from rig.netlist import Net
from rig.place_and_route.constraints import ReserveResourceConstraint
from rig.place_and_route import place
from rig.place_and_route import allocate
from rig.place_and_route import route
from rig_par_diagram import Diagram
from rig_par_diagram.style import Style
from six import iteritems, itervalues
import seaborn as sns


def make_cconv_net(max_cols, max_rows):
    """Make a circular convolution net, save a DOT file description of the
    network, a representative place-and-route diagram and return an indication
    of the traffic density.
    """
    # Compute the number of matrix decompositions
    assert 512 < max_cols or 512 % max_cols == 0
    assert 512 < max_rows or 512 % max_rows == 0

    in_col_divs = 512 // max_cols + (1 if 512 < max_cols else 0)
    in_row_divs = (1028 // max_rows) + (1 if 1028 % max_rows else 0)
    out_col_divs = (1028 // max_cols) + (1 if 1028 % max_cols else 0)
    out_row_divs = 512 // max_rows + (1 if 512 < max_rows else 0)
    print(max_cols, max_rows, in_col_divs, in_row_divs,
          out_col_divs, out_row_divs)

    # Create the original ensemble-arrays
    input_a = [object() for _ in range(32)]
    input_b = [object() for _ in range(32)]

    # Create the input passthrough Node(s)
    in_ptn_a = [[object() for _ in range(in_col_divs)]
                for _ in range(in_row_divs)]
    in_ptn_b = [[object() for _ in range(in_col_divs)]
                for _ in range(in_row_divs)]

    # Create the ensembles in the product net
    ens = [object() for _ in range(2056)]

    # Create the output passthrough Node
    out_ptn = [[object() for _ in range(out_col_divs)]
               for _ in range(out_row_divs)]

    # Create the output ensemble array
    output_ens = [object() for _ in range(32)]

    # Create the nets
    nets = list()

    # Input connections
    in_col_slices = divide_slice(slice(0, len(input_a)),  in_col_divs)
    for i, cols in enumerate(in_col_slices):
        in_row_slices = divide_slice(slice(0, len(ens)),  in_row_divs)

        for j, rows in enumerate(in_row_slices):
            for ptns, ins in [(in_ptn_a, input_a), (in_ptn_b, input_b)]:
                ptn = ptns[j][i]
                nets.extend(Net(e, ptn, weight=16) for e in ins[cols])
                nets.extend(Net(ptn, e, weight=1) for e in ens[rows])

    # Output connections
    cols_slices = list(divide_slice(slice(0, 1028), out_col_divs))
    for i, cols in enumerate(cols_slices):
        col_indices = list(flatten((x, x+1) for
                           x in range(cols.start, cols.stop)))
        out_row_slices = divide_slice(slice(0, len(output_ens)), out_row_divs)

        for j, rows in enumerate(out_row_slices):
            ptn = out_ptn[j][i]

            # Construct the Ensemble -> passthrough Node nets
            # COLUMN ONLY!
            indices = flatten((x, x+1) for x in col_indices)
            nets.extend(Net(ens[e], ptn, weight=1) for e in indices)

            # Construct the passthrough Node -> output nets
            # ROW ONLY!
            nets.extend(Net(ptn, e, weight=16) for e in output_ens[rows])

    # Create the vertices resources
    in_ptn = list(flatten([in_ptn_a, in_ptn_b]))
    vertices_resources = collections.OrderedDict()
    for vx in chain(input_a, input_b, ens, output_ens,
                    in_ptn, flatten(out_ptn)):
        vertices_resources[vx] = {Cores: 1}

    # Create the graph representation
    all_nodes = dict()
    all_nodes.update({a: 'A{}'.format(i) for i, a in enumerate(input_a)})
    all_nodes.update({b: 'B{}'.format(i) for i, b in enumerate(input_b)})
    all_nodes.update({ptn: 'Acconv[{},{}]'.format(i, j) for i, cols in
                      enumerate(in_ptn_a) for j, ptn in enumerate(cols)})
    all_nodes.update({ptn: 'Bcconv[{},{}]'.format(i, j) for i, cols in
                      enumerate(in_ptn_b) for j, ptn in enumerate(cols)})
    all_nodes.update({e: 'Cconv{}'.format(i) for i, e in enumerate(ens)})
    all_nodes.update({ptn: 'CconvC[{}, {}]'.format(i, j) for i, cols in
                      enumerate(out_ptn) for j, ptn in enumerate(cols)})
    all_nodes.update({c: 'C{}'.format(i) for i, c in enumerate(output_ens)})

    g = pgv.AGraph(directed=True)
    g.add_nodes_from(itervalues(all_nodes))

    for n in itervalues(all_nodes):
        g.get_node(n).attr['label'] = ""
        g.get_node(n).attr['shape'] = "point"

    for net in nets:
        for sink in net.sinks:
            so = all_nodes[net.source]
            si = all_nodes[sink]

            g.add_edge(so, si)
            g.get_edge(so, si).attr['weight'] = net.weight

    # Create the core and net styles
    core_style = Style(fill=(0.0, 0.0, 0.0, 0.0),
                       stroke=(0.0, 0.0, 0.0, 0.0),
                       line_width=0.005)
    net_style = Style(stroke=(0.0, 0.0, 0.0, 0.0))

    cols = sns.color_palette("deep", n_colors=4)
    for col, vxs in zip(
            [cols[0],
             cols[1],
             (0.75, ) * 3,
             (0.0, ) * 3,
             cols[2]],
            [flatten([input_a, in_ptn_a]),
             flatten([input_b, in_ptn_b]),
             ens, output_ens, flatten(out_ptn)]):
        # Format the colour
        html_col = "#" + "".join("{:02x}".format(int(c * 255)) for c in col)
        col = col + (1.0, )

        for vx in vxs:
            core_style.set(vx, "fill", col)

            vn = all_nodes[vx]
            g.get_node(vn).attr["color"] = html_col

            for net in nets:
                if net.source is vx:
                    net_style.set(net, "stroke", col)

                    for sink in net.sinks:
                        un = all_nodes[sink]
                        g.get_edge(vn, un).attr["color"] = html_col + "7f"

    # Draw the functional representation
    print("Writing DOT representation...")
    g.write("net_{}_{}.dot".format(max_cols, max_rows))

    # Create the machine
    machine = Machine(24, 12)

    # Place and route, do this a few times
    exp_packets_all = np.zeros((10, machine.height, machine.width))

    for exp_packets in exp_packets_all:
        constraints = [ReserveResourceConstraint(Cores, slice(0, 1))]
        print("Placing...")
        rng = random.Random()
        rng.seed(2)
        placements = place(vertices_resources, nets, machine,
                           constraints, random=rng)
        allocations = allocate(vertices_resources, nets, machine,
                               constraints, placements)
        print("Routing...")
        routes = route(vertices_resources, nets, machine, constraints,
                       placements, allocations)

        # Compute the number of packets to expect at each chip per timestep
        print("Extracting traffic...")
        for net, tree in iteritems(routes):
            # Traverse the tree adding the net weight at every visited chip
            unvisited = collections.deque([tree])
            while unvisited:
                node = unvisited.pop()

                # Add the packet count
                x, y = node.chip
                exp_packets[y, x] += net.weight

                # Traverse onward
                for r, tree in node.children:
                    if r is not None and r.is_link:
                        unvisited.append(tree)

    exp_packets_all *= 1e3  # Convert to packets per second
    exp_packets_all *= 1e-6  # To millions of packets per second

    print("Drawing...")
    # Create the diagram
    d = Diagram(machine, vertices_resources, nets, constraints, placements,
                allocations, routes, core_style=core_style,
                net_style=net_style)

    # Calculate height and width
    width = 4000
    x1, y1, x2, y2 = d.bbox
    w = x2 - x1
    h = y2 - y1
    ratio = h / w

    if ratio < 1.0:
        height = int(width * ratio)
    else:
        height, width = width, int(width / ratio)

    # Create the Cairo context and draw
    surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, width, height)
    ctx = cairo.Context(surface)
    d.draw(ctx, width, height)
    surface.write_to_png("cconv_{}_{}.png".format(max_cols, max_rows))

    return exp_packets_all


if __name__ == "__main__":
    dp = (2056, 1024, 512, 256, 128, 64, 32)
    max_cols_rows = np.array([[a, b] for a in dp for b in dp], dtype=np.int)
    traffic = list()

    for max_cols, max_rows in max_cols_rows:
        traffic.append(make_cconv_net(max_cols, max_rows))

    np.savez_compressed(
        "traffic.npz", max_cols_rows=max_cols_rows, traffic=traffic
    )