arrange.py
# Pranav Gokhale
# Usage: python arrange.py input_trace_file.tr [width height]
# (if width and height are not specified, default to sidelength
# of smallest square big enough to contain all nodes)
# Input file: .tr file
# Output:
# - placement of nodes on the rectangle
import os, sys, math, re
from os.path import basename
from os.path import splitext
module_num_nodes = 0
def main():
trace_to_graph(sys.argv[1])
mapping_dic = {}
nodes_dic, tl_dic, br_dic = parser() # tl = top left coordinate, br = bottom right coordinate
benchname = splitext(basename(sys.argv[1]))[0] + '.p.' + sys.argv[2] + '.yx.' + sys.argv[3] + '.drop.' + sys.argv[4]
for module_name, nodes in nodes_dic.iteritems():
locations = get_locations(nodes, tl=tl_dic[module_name], br=br_dic[module_name])
os.system('rm -f '+benchname+'.delete_*') # partition function creates many files prefixed with delete_
node_map = {v: k for k, v in locations.items()}
mapping = ''
for y in range(br_dic[module_name].y + 1):
for x in range(br_dic[module_name].x + 1):
if node_map.get(Point(x, y)):
mapping += str(node_map.get(Point(x, y)).index) + ' '
else:
mapping += ' '
mapping += '\n'
mapping_dic[module_name] = mapping
replace(mapping_dic, sys.argv[1])
os.system('rm -f '+benchname+'.delete_*')
def trace_to_graph(infile):
"""Converts .tr file to .graph file."""
assert '.tr' in infile
outfile = infile.replace('.tr', '.graph')
infile = open(infile)
outfile = open(outfile, 'w')
for line in infile:
if line.startswith('module:'):
tokens = line.split(' ')
module_name = tokens[1]
outfile.write('module: %s' % module_name)
continue
elif line.startswith('num_nodes:'):
tokens = line.split(' ')
num_nodes = int(tokens[1])
outfile.write('%s\n' % num_nodes)
continue
elif not line.startswith('ID:'): # skip the header lines
continue
elif 'DST:' not in line: # skip single qubit gates
continue
tokens = line.split(' ')
assert tokens[4].startswith('SRC:')
assert tokens[6].startswith('DST:')
src = int(tokens[5])
dst = int(tokens[7])
assert src < num_nodes
assert dst < num_nodes
outfile.write('%s %s\n' % (src, dst))
def parser():
"""Parses nodes from .graph input file. Returns nodes, br, tl.
Example .graph file:
5
0 1
1 2
2 3
3 4
creates
nodes = [
Node(0, [0, 1, 0, 0, 0])
Node(1, [1, 0, 1, 0, 0])
Node(2, [0, 1, 0, 1, 0])
Node(3, [0, 0, 1, 0, 1])
Node(4, [0, 0, 0, 1, 0])
"""
global module_num_nodes
nodes_dic = {}
tl = {}
br = {}
def _square_size(module_num_nodes):
"""Length of the smallest integer sized square with area of at least module_num_nodes.
(size - 1) * (size - 1) < module_num_nodes <= size * size"""
return int(math.ceil(math.sqrt(module_num_nodes)))
def _add_edge(src, dst):
"""Add an undirected edge."""
module_nodes[src].weights[dst] += 1
module_nodes[dst].weights[src] += 1
f = open(sys.argv[1].replace('.tr', '.graph'))
line = f.readline()
module_name = ''
while line != '':
# head of a module
if line.startswith('module:'):
if module_name != '': # save previous module's info
nodes_dic[module_name] = module_nodes
tl[module_name] = Point(0, 0)
br[module_name] = Point(module_width - 1, module_height - 1)
tokens = line.split(' ')
module_name = tokens[1]
module_num_nodes = int(f.readline())
module_nodes = []
if len(sys.argv) == 4: # if width and height are specified from command line
module_width = int(sys.argv[5])
module_height = int(sys.argv[6])
assert module_width * module_height >= module_num_nodes
else:
square_size = _square_size(module_num_nodes)
module_width, module_height = square_size, square_size
for i in range(module_num_nodes):
n = Node(i, module_num_nodes * [0])
module_nodes.append(n)
# body of a module
else:
tokens = line.split(' ') # expected format of each line is "src dst"
assert len(tokens) == 2
src, dst = int(tokens[0]), int(tokens[1])
_add_edge(src, dst)
line = f.readline()
nodes_dic[module_name] = module_nodes # save last module's info
tl[module_name] = Point(0, 0)
br[module_name] = Point(module_width - 1, module_height - 1)
return nodes_dic, tl, br
class Node(object):
def __init__(self, index, weights):
self.index = index
self.weights = weights
def __repr__(self):
return '%s' % (self.index)
class Point(object):
# direction of coordinate system is as follow:
# 0,0 1,0 2,0
# 0,1 1,1 2,1
# 0,2 1,2 2,2
def __init__(self, x, y):
self.x = x
self.y = y
def __repr__(self):
return '(%s, %s)' % (self.x, self.y)
def __hash__(self):
return hash((self.x, self.y))
def __eq__(self, another):
return self.x == another.x and self.y == another.y
def get_locations(nodes, tl, br):
"""Returns a mapping from each node in nodes to a location in the tl<->br square.
tl and br are tuples representing (x, y) coordinates of the top left and bottom right.
"""
# Base cases:
if len(nodes) == 1: # for singleton, only choice is to place in the single spot in 1x1 square
return {nodes[0]: tl}
if len(nodes) == 2: # for two nodes, arbitrarily chose to place the first node in top left
return {nodes[0]: tl, nodes[1]: br}
# Recursive case, need to create and solve subproblems:
ret = {}
num_edges = count_num_edges(nodes)
if num_edges == 0: # for empty graphs, no need to run METIS, just assign arbitrarily
i = 0
for x in range(tl.x, br.x+1):
for y in range(tl.y, br.y+1):
if i < len(nodes):
ret.update({nodes[i]: Point(x,y)})
i += 1
return ret
filename = splitext(basename(sys.argv[1]))[0] + '.p.' + sys.argv[2] + '.yx.' + sys.argv[3] + '.drop.' + sys.argv[4] + '.' +\
'_'.join(['delete', str(tl.x), str(tl.y), str(br.x), str(br.y)])
# special case for the very first call of get_locations. For example, suppose that there are
# 97 nodes on a 10x10 grid. Instead of dividing the 97 nodes into 2 equal partitions, we should
# divide them into a partition of 90 nodes and a partition of 7 nodes. The former should be
# placed on a 10x9 grid and te latter should be placed on a 1x7 grid.
if len(nodes) < (br.x - tl.x + 1) * (br.y - tl.y + 1):
assert tl == Point(0, 0)
size_tl_nodes = (br.x + 1) * int(len(nodes) / (br.x + 1))
if size_tl_nodes == len(nodes):
ret.update(get_locations(nodes, tl=Point(0, 0), br=Point(br.x, len(nodes) / (br.x + 1) - 1)))
return ret
nodes_tl, nodes_br = partition(nodes, size_tl_nodes, filename)
# complicated indexing here. As an example, for the 97 into 10x10 case, we want to send 90 nodes
# to a rectangle spanned by tl=Point(0, 0) and br=Point(9, 8) and we want to send 7 nodes to a
# rectangle spanned by tl=Point(0, 9) and br=Point(6, 9)
ret.update(get_locations(nodes_tl, tl=Point(0, 0), br=Point(br.x, len(nodes) / (br.x + 1) - 1)))
ret.update(get_locations(nodes_br, tl=Point(0, len(nodes) / (br.x + 1)), br=Point(len(nodes) % (br.x + 1) - 1, len(nodes) / (br.x + 1))))
return ret
if br.x - tl.x > br.y - tl.y: # if rectangle is wider than tall, split on y axis
half = tl.x + (br.x - tl.x - 1) / 2
size_tl_nodes = (half - tl.x + 1) * (br.y - tl.y + 1)
else: # split on x axis
half = tl.y + (br.y - tl.y - 1) / 2
size_tl_nodes = (br.x - tl.x + 1) * (half - tl.y + 1)
nodes_tl, nodes_br = partition(nodes, size_tl_nodes, filename)
if br.x - tl.x > br.y - tl.y: # if rectangle is wider than tall, split on y axis
ret.update(get_locations(nodes_tl, tl=tl, br=Point(half, br.y)))
ret.update(get_locations(nodes_br, tl=Point(half + 1,tl.y), br=br))
else: # split on x axis
ret.update(get_locations(nodes_tl, tl=tl, br=Point(br.x, half)))
ret.update(get_locations(nodes_br, tl=Point(tl.x, half + 1), br=br))
return ret
def count_num_edges(nodes):
num_edges = 0
for src_node in nodes:
for dst_node in nodes:
if src_node.weights[dst_node.index]:
num_edges += 1
num_edges /= 2 # don't double count edges
return num_edges
def partition(nodes, size_nodes_tl, filename):
"""Returns nodes_tl, nodes_br."""
weight_0 = float(size_nodes_tl) / len(nodes)
num_edges = count_num_edges(nodes)
f = open(filename, 'w')
f.write('%s %s 001\n' % (len(nodes), num_edges))
for src_node in nodes:
x = 1
for dst_node in nodes:
if src_node.weights[dst_node.index]:
f.write('%s %s ' % (x, src_node.weights[dst_node.index]))
x += 1
f.write('\n')
f.close()
benchname = splitext(basename(sys.argv[1]))[0] + '.p.' + sys.argv[2] + '.yx.' + sys.argv[3] + '.drop.' + sys.argv[4]
f = open(benchname+'.delete_tpwgts', 'w') # target partition weights file
f.write('0 = %s' % str(weight_0))
f.close()
os.system('gpmetis -ptype=rb -tpwgts='+benchname+'.delete_tpwgts -ufactor=1 -ncuts=500 ' + filename + ' 2 > /dev/null')
f = open(filename + '.part.2')
nodes_tl, nodes_br = [], []
index = 0
for line in f:
if line == '0\n':
nodes_tl.append(nodes[index])
else:
nodes_br.append(nodes[index])
index += 1
assert len(nodes_tl) + len(nodes_br) == len(nodes), 'nodes went missing somehow :('
# METIS is not perfect and will not always actually yield size_nodes_tl nodes in nodes_tl
# in this case, just pop enough nodes from one list to the other, until sizes are appropriate
# TODO: pick the nodes to transfer between lists more intelligently
while len(nodes_tl) > size_nodes_tl:
nodes_br.append(nodes_tl.pop())
while size_nodes_tl > len(nodes_tl):
nodes_tl.append(nodes_br.pop())
return nodes_tl, nodes_br
def replace(mapping_dic, f_tr):
"""Replace the nodes in the .tr file using the mapping and generate .opt.tr file."""
f_opt_tr = f_tr.replace('.tr', '.opt.tr')
f_tr = open(f_tr)
f_opt_tr = open(f_opt_tr, 'w')
for line in f_tr:
if line.startswith('module: '):
tokens = line.split(' ')
module_name = tokens[1]
mapping = mapping_dic[module_name].split()
d = {}
i = 0
for item in mapping:
d[int(item)] = i
i += 1
if 'SRC' not in line:
f_opt_tr.write(line)
continue
line = re.sub(r'SRC: (\d+)', lambda match: 'SRC: '+str(d.get(int(match.group(1)), 'SOMETHING WENT WRONG' + match.group(1))), line)
line = re.sub(r'DST: (\d+)', lambda match: 'DST: '+str(d.get(int(match.group(1)), 'SOMETHING WENT WRONG' + match.group(1))), line)
f_opt_tr.write(line)
if __name__ == "__main__":
main()
