report.py

import json
import argparse
import os
import re
import sys
import csv
from collections import defaultdict

def load_benchmark_data(filename):
    try:
        with open(filename, 'r') as f:
            return json.load(f)
    except Exception as e:
        print(f"Error loading {filename}: {str(e)}")
        return None

def normalize_test_name(name):
    name = re.sub(r'/iterations:\d+', '', name)
    name = re.sub(r'/repeats:\d+', '', name)
    
    if name.endswith('/'):
        name = name[:-1]
    
    if name.endswith('_mean'):
        name = name[:-5]
    
    return name

def extract_function_name(test_case_name):
    base_name = test_case_name.split('/')[0] if '/' in test_case_name else test_case_name
    return base_name.lower()

def remove_common_prefix(filenames):

    if not filenames:
        return filenames
    
    base_names = [os.path.splitext(os.path.basename(f))[0] for f in filenames]
    
    prefix = os.path.commonprefix(base_names)
    
    if len(prefix) > 3:
        return [name[len(prefix):] for name in base_names]
    return base_names

def calculate_group_averages(file_data, all_cases):
    groups = defaultdict(list)
    for case in all_cases:
        func_name = extract_function_name(case)
        groups[func_name].append(case)
    
    group_results = {}
    for func_name, cases in groups.items():
        group_times = [[] for _ in file_data]
        
        speedups = [[] for _ in range(len(file_data) - 1)]
        
        for case in cases:
            times = []
            for data in file_data:
                time_val = data['means'].get(case)
                times.append(time_val)
                
            for i, time_val in enumerate(times):
                if time_val is not None:
                    group_times[i].append(time_val)
            
            base_time = times[0]
            if base_time is not None and base_time > 0:
                for i in range(1, len(times)):
                    if times[i] is not None and times[i] > 0:
                        speedup = base_time / times[i]
                        speedups[i - 1].append(speedup)
        
        avg_times = []
        for times in group_times:
            if times:
                avg_times.append(sum(times) / len(times))
            else:
                avg_times.append(None)
        
        avg_speedups = []
        for speedup_list in speedups:
            if speedup_list:
                avg_speedups.append(sum(speedup_list) / len(speedup_list))
            else:
                avg_speedups.append(None)
        
        group_results[func_name] = {
            'avg_times': avg_times,
            'avg_speedups': avg_speedups,            'case_count': len(cases)
        }
    
    return group_results

def print_detailed_results(file_data, all_cases, args, csv_rows=None):
    headers = ["Test Case"]
    
    if not args.speedup_only or len(file_data) == 1:
        for data in file_data:
            headers.append(data['display_name'])
    
    if len(file_data) > 1:
        for i in range(1, len(file_data)):
            headers.append(f"{file_data[i]['display_name']}/Speedup")
    
    header_line = ",".join(headers)
    print(header_line)
    
    for case in sorted(all_cases):
        row = [case]
        
        times = []
        for data in file_data:
            time_val = data['means'].get(case)
            times.append(time_val)
        
        if not args.speedup_only or len(times) == 1:
            for t in times:
                value_str = f"{t:.4f}" if t is not None else "N/A"
                row.append(value_str)
        
        if len(times) > 1:
            base_time = times[0]
            if base_time is None or base_time == 0:
                for i in range(1, len(times)):
                    row.append("N/A")
            else:
                for i in range(1, len(times)):
                    if times[i] is not None and times[i] != 0:
                        speedup = base_time / times[i]
                        row.append(f"{speedup:.4f}")
                    else:
                        row.append("N/A")
        
        print(",".join(row))
        
        if csv_rows is not None:
            csv_rows.append(row)
    
    return headers

def print_group_averages(file_data, group_results, args, csv_rows=None):
    headers = ["Function"]
    
    if not args.speedup_only or len(file_data) == 1:
        for data in file_data:
            headers.append(f"{data['display_name']}/Avg")
    
    if len(file_data) > 1:
        for i in range(1, len(file_data)):
            headers.append(f"Speedup/{file_data[i]['display_name']}/Avg")
    
    header_line = ",".join(headers)
    print(header_line)
    
    for func_name in sorted(group_results.keys()):
        result = group_results[func_name]
        row = [func_name]
        
        avg_times = result['avg_times']
        avg_speedups = result['avg_speedups']        
        if not args.speedup_only or len(avg_times) == 1:
            for t in avg_times:
                value_str = f"{t:.4f}" if t is not None else "N/A"
                row.append(value_str)

        if len(file_data) > 1:
            for speedup in avg_speedups:
                if speedup is not None:
                    row.append(f"{speedup:.4f}")
                else:
                    row.append("N/A")
        print(",".join(row))
        
        if csv_rows is not None:
            csv_rows.append(row)
    
    return headers

def main():
    parser = argparse.ArgumentParser(description='Compare Google Benchmark JSON results')
    parser.add_argument('json_files', nargs='+', help='JSON benchmark files to compare')
    parser.add_argument('--speedup-only', action='store_true',
                        help='Only show speedup ratios (omit time values)')
    parser.add_argument('--output', type=str, default=None,
                        help='Save CSV output to specified file')
    parser.add_argument('--group', action='store_true',
                        help='Group by function name and show average speedup')
    parser.add_argument('--filter', type=str, default=None,
                        help='Filter by function name (comma-separated list, e.g., "relu,convsamp")')
    args = parser.parse_args()

    if len(args.json_files) < 1:
        print("Error: At least one JSON file is required", file=sys.stderr)
        return

    display_names = remove_common_prefix(args.json_files)
    file_data = []
    for i, file_path in enumerate(args.json_files):
        if not os.path.exists(file_path):
            print(f"Error: File not found - {file_path}", file=sys.stderr)
            return
            
        data = load_benchmark_data(file_path)
        if not data:
            return
            
        mean_times = parse_mean_times(data)
        
        if not mean_times:
            print(f"Error: No benchmark results found in {file_path}", file=sys.stderr)
            return
            
        file_data.append({
            'path': file_path,
            'display_name': display_names[i],
            'means': mean_times
        })

    all_cases = set()
    for data in file_data:
        all_cases.update(data['means'].keys())
    
    if args.filter:
        filter_funcs = set(f.strip().lower() for f in args.filter.split(','))
        filtered_cases = set()
        for case in all_cases:
            func_name = extract_function_name(case)
            if func_name in filter_funcs:
                filtered_cases.add(case)
        
        if not filtered_cases:
            print(f"Warning: No test cases match filter '{args.filter}'", file=sys.stderr)
            print(f"Available functions: {', '.join(sorted(set(extract_function_name(c) for c in all_cases)))}", file=sys.stderr)
            return
        
        all_cases = filtered_cases
    
    csv_rows = [] if args.output else None
    
    if args.group:
        group_results = calculate_group_averages(file_data, all_cases)
        headers = print_group_averages(file_data, group_results, args, csv_rows)
    else:
        headers = print_detailed_results(file_data, all_cases, args, csv_rows)
    
    if args.output and csv_rows:
        try:
            with open(args.output, 'w', newline='') as csv_file:
                writer = csv.writer(csv_file)
                
                writer.writerow(headers)
                
                for row in csv_rows:
                    writer.writerow(row)
                    
            print(f"\nCSV results saved to {args.output}", file=sys.stderr)
        except Exception as e:
            print(f"Error saving CSV to {args.output}: {str(e)}", file=sys.stderr)

def parse_mean_times(data):
    means = {}
    for item in data.get('benchmarks', []):
        if (item.get('run_type') == 'aggregate' and 
            item.get('aggregate_name') == 'mean'):
            normalized_name = normalize_test_name(item['name'])
            means[normalized_name] = item['cpu_time']
    return means

if __name__ == "__main__":
    main()