03_atpg.py

import os 
import glob 
from torch_geometric.data import Data, InMemoryDataset
import deepgate as dg 
import torch
import random
import time
import threading
import copy
import numpy as np 
import re

from utils.utils import run_command, hash_arr
from parse_graph import parse_sdf
import utils.circuit_utils as circuit_utils

v_dir = './data/raw_v'
atpg_dir = './data/atpg'
abccrc_path = '/Users/zhengyuanshi/opt/abc/abc.rc'
lib_path = 'lib/sky130_fd_sc_hd__ff_n40C_1v95.lib'

RESYN = False

if __name__ == '__main__':
    v_list = glob.glob(os.path.join(v_dir, '*/*.v'))
    tot_time = 0
    no_succ = 0
    failed_list = []
    if not os.path.exists(atpg_dir):
        os.makedirs(atpg_dir)
    
    for v_k, v_path in enumerate(v_list):
        
        start_time = time.time()
        arr = v_path.replace('.v', '').split('/')
        design_name = arr[-2]
        module_name = arr[-1]
        circuit_name = '{}/{}'.format(design_name, module_name)
        if not os.path.exists(v_path):
            print('[ERROR] File not found: {}'.format(v_path))
            continue
        print('[INFO] Parse: {}'.format(v_path))
        
        design_aig_dir = os.path.join(atpg_dir, design_name)
        if not os.path.exists(design_aig_dir):
            os.makedirs(design_aig_dir)
        atpg_aig_path = os.path.join(design_aig_dir, '{}.aig'.format(module_name))
        atpg_v_path = os.path.join(design_aig_dir, '{}.v'.format(module_name))
        
        # Parse verilog to construct ATPG case 
        f = open(v_path, 'r')
        lines = f.readlines()
        f.close()
        f = open(atpg_v_path, 'w')
        if len(lines) < 1000:
            continue
        
        for line in lines:
            if 'endmodule' in line:
                continue
            elif 'module' in line:
                f.write(line)
                f.write('output po;\n')
            elif line.lstrip()[:5] == 'input':
                input_arr = line.replace('input', '').replace(';', '').replace('\n', '').replace(' ', '').split(',')
                pi_list = []
                for input_name in input_arr:
                    if len(input_name) > 0:
                        pi_list.append(input_name)
                new_line = 'input '
                for k, input_name in enumerate(pi_list):
                    if k == len(pi_list) - 1:
                        new_line += '{};\n'.format(input_name)
                    else:
                        new_line += '{},'.format(input_name)
                f.write(new_line)
            elif line.lstrip()[:6] == 'output':
                output_arr = line.replace('output', '').replace(';', '').replace('\n', '').replace(' ', '').split(',')
                po_list = []
                for output_name in output_arr:
                    if len(output_name) > 0:
                        po_list.append(output_name)
                new_line = 'wire '
                for k, output_name in enumerate(po_list):
                    if k == len(po_list) - 1:
                        new_line += '{};\n'.format(output_name)
                    else:
                        new_line += '{},'.format(output_name)
                f.write(new_line)
                new_line = 'wire '
                for k, output_name in enumerate(po_list):
                    if k == len(po_list) - 1:
                        new_line += '{}_r;\n'.format(output_name)
                    else:
                        new_line += '{}_r,'.format(output_name) 
                f.write(new_line)                   
                new_line = 'wire '
                for k, output_name in enumerate(po_list):
                    if k == len(po_list) - 1:
                        new_line += '{}_xor;\n'.format(output_name)
                    else:
                        new_line += '{}_xor,'.format(output_name) 
                f.write(new_line)                   
            elif line.lstrip()[:4] == 'wire':
                wire_arr = line.replace('wire', '').replace(';', '').replace('\n', '').replace(' ', '').split(',')
                wire_list = []
                for wire_name in wire_arr:
                    if len(wire_name) > 0:
                        wire_list.append(wire_name)
                new_line = 'wire '
                for k, wire_name in enumerate(wire_list):
                    if k == len(wire_list) - 1:
                        new_line += '{};\n'.format(wire_name)
                    else:
                        new_line += '{},'.format(wire_name)
                f.write(new_line)
                new_line = 'wire '
                for k, wire_name in enumerate(wire_list):
                    if k == len(wire_list) - 1:
                        new_line += '{}_r;\n'.format(wire_name)
                    else:
                        new_line += '{}_r,'.format(wire_name)
                f.write(new_line)
            elif ';' not in line:
                f.write(line)
            else:
                arr = line.split(' ')
                cell_type = arr[0]
                cell_name = arr[1]
                new_line = copy.deepcopy(line)
                for pin in wire_list + po_list:
                    new_line = re.sub(r'\b' + re.escape(pin) + r'\b', f'{pin}_r', new_line)
                new_line = re.sub(r'(\w+_Cell)\b', r'\1_r', new_line)
                f.write(line)
                f.write(new_line)
        xor_list = []
        for po in po_list:
            new_line = 'sky130_fd_sc_hd__xor2_1 {}_xor_Cell (.A({}), .B({}_r), .X({}_xor));\n'.format(
                po, po, po, po
            )
            f.write(new_line)
            xor_list.append('{}_xor'.format(po))
        no_or_gate = 0
        while len(xor_list) > 4:
            or_gate = 'or_{}'.format(no_or_gate)
            no_or_gate += 1
            new_line = 'wire {};\n'.format(or_gate)
            f.write(new_line)
            new_line = 'sky130_fd_sc_hd__or4_1 {}_Cell (.A({}), .B({}), .C({}), .D({}), .X({}));\n'.format(
                or_gate, xor_list[0], xor_list[1], xor_list[2], xor_list[3], or_gate
            )
            f.write(new_line)
            xor_list = xor_list[4:]
            xor_list.append(or_gate)
        if len(xor_list) == 4:
            new_line = 'sky130_fd_sc_hd__or4_1 po_Cell (.A({}), .B({}), .C({}), .D({}), .X(po));\n'.format(
                xor_list[0], xor_list[1], xor_list[2], xor_list[3]
            )
            f.write(new_line)
        elif len(xor_list) == 3:
            new_line = 'sky130_fd_sc_hd__or3_1 po_Cell (.A({}), .B({}), .C({}), .X(po));\n'.format(
                xor_list[0], xor_list[1], xor_list[2]
            )
            f.write(new_line)
        elif len(xor_list) == 2:
            new_line = 'sky130_fd_sc_hd__or2_1 po_Cell (.A({}), .B({}), .X(po));\n'.format(
                xor_list[0], xor_list[1]
            )
            f.write(new_line)
            
        f.write('endmodule\n')
        f.close()
        
        print(atpg_v_path)
        
        
        for k in range(int(min(len(wire_list) * 0.01, 1000))):
            wire = random.choice(wire_list)
            f = open(atpg_v_path, 'r')
            lines = f.readlines()
            f.close()
            for const in ['0', '1']:
                tmp_atpg_v_path = atpg_v_path.replace('.v', '_tmp.v')
                f = open(tmp_atpg_v_path, 'w')
                for line in lines:
                    if '{}_Cell '.format(wire) in line:
                        new_line = '_const_{}_ {}_const_{}_Cell (.z({}));\n'.format(const, wire, const, wire)
                        f.write(new_line)
                    else:
                        f.write(line)
                f.close()
                abc_cmd = 'abc -c \"read {}; read_verilog -m {}; strash; sat -C 100000;\"'.format(lib_path, tmp_atpg_v_path)
                stdout, _ = run_command(abc_cmd)
                arr = stdout[-2].replace('\n', '').replace(' ', '').replace('sec', '').replace('Time', '').split('=')
                if len(arr) != 2:
                    continue
                res = arr[0]
                solve_time = float(arr[1])
                # Find hard case
                if res != 'SATISFIABLE' and res != 'UNSATISFIABLE':
                    tmp_atpg_aig_path = atpg_aig_path.replace('.aig', '_{}_{}.aig'.format(wire, const))
                    abc_cmd = 'abc -c \"read {}; read_verilog -m {}; strash; write_aiger {};\"'.format(lib_path, tmp_atpg_v_path, tmp_atpg_aig_path)
                    stdout, _ = run_command(abc_cmd)
                    print('Save: {}'.format(tmp_atpg_aig_path))
                if os.path.exists(tmp_atpg_v_path):
                    os.remove(tmp_atpg_v_path)
                print('Wire: {}, Const: {}, Res: {}, Time: {:.2f}s'.format(wire, const, res, solve_time))
            
                        
        
        
        # Statistics
        print('Parse: {} ({:} / {:}), Time: {:.2f}s, ETA: {:.2f}s'.format(
            circuit_name, v_k, len(v_list), 
            tot_time, tot_time / ((v_k + 1) / len(v_list)) - tot_time
        ))
        tot_time += time.time() - start_time
    
    
    print()
    print('='*20)
    failed_path = './aig_to_aig_failed.txt'
    f = open(failed_path, 'w')
    for failed in failed_list:
        f.write('{}\n'.format(failed))
        print(failed)
    f.close()
    print()
    
    print('Total: {}, Succ: {}'.format(len(v_list), no_succ))