datapath.py

import time
import re
from dataclasses import dataclass

# Defined locally
from utils import run, grep_stat

################################################################################
# Data Collection Classes
################################################################################
# @dataclass
class DatapathDataSet:
    def __init__(self, name, dir, output_dir, bw=None):
        self.name = name
        self.dir = dir
        self.output_dir = output_dir
        self.bw = bw
        self.stats = {}
        self.sv_file = f"benchmarks/{self.dir}/sv/{self.dir}.sv"
        self.comb_mlir_file = f"{self.output_dir}/{self.dir}.comb.mlir"
        self.mlir_aig_file = f"{self.output_dir}/{self.dir}.{self.name}.aig.mlir"
        self.aiger_file = f"{self.output_dir}/{self.dir}.{self.name}.aig"
        self.smt2_file = f"{self.output_dir}/{self.dir}.{self.name}.smt2"
        # self.smt2_file = f"benchmarks/{self.dir}/smtlib/{self.dir}.circt.{self.bw}_bit.smt2"

    ############################################################################
    # Front-End Synthesis
    ############################################################################
    # Convert sv to comb mlir and run circt-synth
    def run_circt_synth(self, options=""):
        # Generate comb MLIR from SV
        run(f'circt-verilog {self.sv_file} -G BW={self.bw} -o {self.comb_mlir_file}')

        # Run once with timing to get longest path
        run(f'circt-synth {self.comb_mlir_file} {options} -o {self.mlir_aig_file} --output-longest-path={self.output_dir}/{self.dir}.{self.name}.path')
        self.stats['circt_levels'] = grep_stat(f"{self.output_dir}/{self.dir}.{self.name}.path", r'Maximum path delay: ([0-9]+)')
        
        # Run again to get clean timing info - without longest path analysis
        # Comb --> AIG in MLIR
        start = time.time()
        run(f'circt-synth {self.comb_mlir_file} {options} -o {self.mlir_aig_file}')
        self.stats['circt_time'] = time.time() - start
        
        # AIG in MLIR --> AIGER format
        start1 = time.time()
        run(f'circt-translate {self.output_dir}/{self.dir}.{self.name}.aig.mlir --export-aiger -o {self.aiger_file}')
        self.stats['mlir_to_aiger_time'] = time.time() - start1
        
    # Convert sv to comb mlir and run circt-synth, then verify equivalence of the two mlir circuits
    def verify_circt_synth(self, options=""):
        # Can only verify circuits expressed in comb - so run with --convert-to-comb    
        run(f'circt-synth {self.comb_mlir_file} {options} --convert-to-comb -o {self.mlir_aig_file}')

        # TODO: implement equivalence checking here

    # Run yosys synthesis on either sv or aiger input from circt-synth flow
    def run_yosys_synth(self, input_fmt):
        
        yosys_cmd = f' synth;'
        # Process System Verilog
        if input_fmt == "sv":
            yosys_cmd = f'read_verilog -sv {self.sv_file}; chparam -set BW {self.bw} {self.dir};' + yosys_cmd

            # Export AIGER
            yosys_cmd += f' aigmap; write_aiger {self.aiger_file};'
            
        # Process AIGER (generated by circt-synth)
        elif input_fmt == "aiger":
            yosys_cmd = f'read_aiger {self.aiger_file};' + yosys_cmd
        else:
            print(f"Error: Unsupported input format '{input_fmt}' for yosys synthesis.")
            assert False

        stat_file = f"{self.output_dir}/{self.dir}.{self.name}.yosys_stat"
        start = time.time()
        # Run Yosys synthesis and generate an AIGER file if processing sv
        run(f'yosys -f verilog -p "{yosys_cmd}" > {stat_file}')
        self.stats['yosys_time'] = time.time() - start
        self.stats['yosys_cells'] = grep_stat(stat_file, r'Number of cells: +([0-9]+)')
    
    ############################################################################
    # Technology Mapping
    ############################################################################
    def run_abc_techmapping(self, area, delay):
        # Run ABC technology mapping on the AIGER file
        start = time.time()
        run(f'abc -c "read_genlib libraries/asap7.genlib; read {self.aiger_file}; strash; map; print_stats" > {self.output_dir}/{self.dir}.{self.name}.abc_stat')
        self.stats['abc_time'] = time.time() - start

        self.stats['abc_area'] = grep_stat(f"{self.output_dir}/{self.dir}.{self.name}.abc_stat", r'area =+([0-9.]+)')
        self.stats['abc_delay'] = grep_stat(f"{self.output_dir}/{self.dir}.{self.name}.abc_stat", r'delay =+([0-9.]+)')
        area[self.name].append(float(self.stats['abc_area']))
        delay[self.name].append(float(self.stats['abc_delay']))

    def run_emap_techmapping(self):
        # Run emap technology mapping on the AIGER file
        start = time.time()
        run(f'emap {self.aiger_file} > {self.output_dir}/{self.dir}.{self.name}.emap_stat')
        self.stats['emap_time'] = time.time() - start

        self.stats['emap_area'] = grep_stat(f"{self.output_dir}/{self.dir}.{self.name}.emap_stat", r'Area ([0-9.]+)')
        self.stats['emap_delay'] = grep_stat(f"{self.output_dir}/{self.dir}.{self.name}.emap_stat", r'Delay ([0-9.]+)')

    ############################################################################
    # Verification
    ############################################################################
    def generate_smtlib(self, options=""):
        # Generate comb MLIR from SV
        run(f'circt-verilog {self.sv_file} -G BW={self.bw} -o {self.comb_mlir_file}')
        run(f'circt-synth {self.comb_mlir_file} {options} --convert-to-comb -o {self.mlir_aig_file}')
        
        run(f'circt-lec {self.comb_mlir_file} {self.mlir_aig_file} --c1 {self.dir} --c2 {self.dir} --emit-smtlib -o {self.smt2_file}')
        run(f'sed -i \'/(reset)/d\' {self.smt2_file}')
    
    def run_z3(self, options=""):
        start = time.time()
        run(f'z3 -T:60 {self.smt2_file}', f'{self.output_dir}/equiv.log')
        self.stats["z3_time"] = time.time() - start
        with open(f'{self.output_dir}/equiv.log', 'r') as f:
            log = f.read()
            match = re.search(r'sat', log)
            if match:
                self.stats["equiv"] = "NEQ"
            
            match = re.search(r'unknown', log)
            if match: 
                self.stats["equiv"] = "UNKNOWN"
            
            match = re.search(r'unsat', log)
            if match:
                self.stats["equiv"] = "EQ"
    
    def print_header(self):
        header = ""
        for key in self.stats.keys():
            header += f" {key} "
        return header.rstrip(' | ')
    
    def print_string(self):
        line = ""
        for key in self.stats.keys():
            width = len(key)
            val = self.stats[key]
            try:
                num = float(val)
                if num == int(num):
                    formatted = str(int(num))
                else:
                    formatted = f"{num:.2f}"
            except (ValueError, TypeError):
                formatted = str(val)
            line += f" {formatted:>{width}} "
        return line.rstrip(' | ')