tool_extract_subcircuit.py

import argparse
import json
import re
import copy

from py_verilog_parser import parse_verilog, generate_verilog
from py_verilog_parser.verilog import Netlist
from circuit_graph import convert_to_graph, convert_to_verilog, Graph

if __name__ == '__main__':
	parser = argparse.ArgumentParser(description='Verilog transform utility')
	parser.add_argument('--input', required=True, dest='input_files', action='append', default=[], help='The Verilog input file')
	parser.add_argument('--output', required=True, dest='output_file', action='store', help='The Verilog output file')
	parser.add_argument('--top', required=True, dest='top_module', action='store', help='The top-level module')
	parser.add_argument('--flipflops', dest='flipflops', action='store', default='.*', help='The regex selecting the relevant flip-flops')
	parser.add_argument('--ports', dest='ports', action='store', default='.*', help='The regex selecting the relevant ports')
	args = parser.parse_args()

	modules = []
	for input_file in args.input_files:
		print(f'Info: Reading {input_file}')
		with open(input_file, 'rt', encoding='utf-8') as input:
			ast = parse_verilog(input.read())
			modules += ast.modules

	modules = dict([(module.module_name, module) for module in modules])
	if args.top_module not in modules:
		print(f'Could not find top-level module {args.top_module}')
		exit(1)

	graph = convert_to_graph(args.top_module, modules)

	regex_flipflops = re.compile(args.flipflops)
	regex_ports = re.compile(args.ports)

	def _has_match(names, regex):
		return any([regex.match(name) for name in names])

	def _has_base_name(names):
		return any([(not '/' in name) for name in names])

	flipflop_nodes = {}
	port_nodes = {}
	for index, node in enumerate(graph.elements):
		if node.type in ['gate'] and node.sub_type in ['flipflop'] \
				and _has_match(node.names, regex_flipflops):
			for name in node.names:
				flipflop_nodes[name] = None
		elif node.type in ['input', 'output'] \
				and _has_match(node.names, regex_ports) \
				and _has_base_name(node.names):
			for name in node.names:
				port_nodes[name] = None

	def next(connections, next_active, visited):
		for connection in connections:
			if connection not in graph.names:
				continue
			next_node = graph.names[connection]
			for name in next_node.names:
				if name not in visited:
					next_active[name] = None

	active = {**port_nodes, **flipflop_nodes}
	visited_forward = {}
	while active:
		next_active = {}
		for name in active:
			node = graph.names[name]
			if node.type in ['input', 'output']:
				for output, connections in node.outputs.items():
					next(connections, next_active, visited_forward)
				if '/' in name:
					parent = name[0:name.index('/')]
					next([parent], next_active, visited_forward)
			elif node.type in ['constant', 'net', 'assignment']:
				for output, connections in node.outputs.items():
					next(connections, next_active, visited_forward)
			elif node.type in ['gate']:
				if node.sub_type in ['flipflop', 'latch'] and not name in visited_forward:
					for output, connections in node.outputs.items():
						next(connections, next_active, visited_forward)
				elif node.sub_type not in ['flipflop', 'latch']:
					for output, connections in node.outputs.items():
						next(connections, next_active, visited_forward)
			elif node.type in ['module']:
				for output, connections in node.outputs.items():
					next(connections, next_active, visited_forward)
			visited_forward[name] = None
		active = next_active

	active = {**port_nodes, **flipflop_nodes}
	visited_backward = {}
	while active:
		next_active = {}
		for name in active:
			node = graph.names[name]
			if node.type in ['input', 'output']:
				for output, connections in node.inputs.items():
					next(connections, next_active, visited_backward)
				if '/' in name:
					parent = name[0:name.index('/')]
					next([parent], next_active, visited_backward)
			elif node.type in ['constant', 'net', 'assignment']:
				for output, connections in node.inputs.items():
					next(connections, next_active, visited_backward)
			elif node.type in ['gate']:
				if node.sub_type in ['flipflop', 'latch'] and not name in visited_backward:
					for output, connections in node.inputs.items():
						next(connections, next_active, visited_backward)
				elif node.sub_type not in ['flipflop', 'latch']:
					for output, connections in node.inputs.items():
						next(connections, next_active, visited_backward)
			elif node.type in ['module']:
				for output, connections in node.inputs.items():
					next(connections, next_active, visited_backward)
			visited_backward[name] = None
		active = next_active

	visited = set(visited_forward.keys()).intersection(set(visited_backward.keys()))

	def _filter_names(element, without_prefix=False):
		return [name for name in element.names if name in visited \
			and (not without_prefix or '/' not in name)]

	def _filter_cells(element):
		return [name for name in element.names if name in visited]

	def _get_shortest_name(names):
		names = sorted([(name.count('/'), name) for name in names])
		_, name = names[-1]
		return name

	def _map_name(element):
		new_element = copy.deepcopy(element)
		new_element.names = [_get_shortest_name(element.names)]
		new_element.inputs = {}
		new_element.outputs = {}
		for port_name, connections in element.inputs.items():
			if isinstance(connections, str) or isinstance(connections, int):
				connections = [connections]
			if any([conn for conn in connections if conn in visited]):
				new_element.inputs[port_name.replace(' [0]', '')] = [_get_shortest_name(connections)]
		for port_name, connections in element.outputs.items():
			if isinstance(connections, str) or isinstance(connections, int):
				connections = [connections]
			if any([conn for conn in connections if conn in visited]):
				new_element.outputs[port_name.replace(' [0]', '')] = [_get_shortest_name(connections)]
		return new_element

	new_graph = Graph()
	new_graph.inputs = [_map_name(port) for port in graph.inputs if _filter_names(port, without_prefix=True)]
	new_graph.outputs = [_map_name(port) for port in graph.outputs if _filter_names(port, without_prefix=True)]
	new_graph.assignments = [_map_name(assignment) for assignment in graph.assignments if _filter_names(assignment)]
	new_graph.wires = [_map_name(wire) for wire in graph.wires if _filter_names(wire)]
	new_graph.gates = []
	new_graph.modules = [_map_name(elem) for elem in graph.modules if _filter_cells(elem)]
	new_graph.elements = new_graph.inputs + new_graph.outputs \
		+ new_graph.assignments + new_graph.wires + new_graph.gates + new_graph.modules

	module = convert_to_verilog(args.top_module, new_graph)
	verilog = generate_verilog(Netlist([module]), [])
	with open(args.output_file, 'wt', encoding='utf-8') as stream:
		stream.write(verilog)