bench_sanger.py

import argparse
from pathlib import Path

import pandas as pd


def bert_base_gflops(seq_len):
    HIDDEN_SIZE = 768
    linear_flops = seq_len * HIDDEN_SIZE * HIDDEN_SIZE * 2 * 3
    qk_flops = seq_len * seq_len * HIDDEN_SIZE * 2
    pv_flops = seq_len * seq_len * HIDDEN_SIZE * 2
    out_proj_flops = seq_len * HIDDEN_SIZE * HIDDEN_SIZE * 2
    stage1_flops = linear_flops
    stage2_flops = qk_flops + pv_flops + out_proj_flops
    stage1_gflops = stage1_flops / 1e9
    stage2_gflops = stage2_flops / 1e9
    return stage1_gflops, stage2_gflops


def calc_sanger_latency(sparsity, load_balance, seq_len):
    PE_ARRAY_SIZE = 64 * 16
    STAGE1_GOPS = PE_ARRAY_SIZE * 1 * 2  # pe-size * 1(GHz) * 2(ops/mac) = 2048
    stage2_gops = STAGE1_GOPS / sparsity * load_balance
    stage1_gflops, stage2_gflops = bert_base_gflops(seq_len)
    stage1_lat = stage1_gflops / STAGE1_GOPS
    stage2_lat = stage2_gflops / stage2_gops
    total_lat = stage1_lat + stage2_lat
    return total_lat


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--sparsity", default=None, type=float, required=False)
    parser.add_argument("--load_balance", default=None, type=float, required=False)
    parser.add_argument("--seq_len", default=512, type=int, required=False)
    parser.add_argument("--csv_file", default="load_balance.csv", type=str, required=False, 
                        help="Path to the csv file generated by gen_sparsity_mask.")
    args = parser.parse_args()
    
    if args.sparsity is not None:
        assert args.load_balance is not None
        total_lat = calc_sanger_latency(args.sparsity, args.load_balance, args.seq_len)
        print(f"Sanger Latency: {total_lat * 1000:.3f} ms")
    else:
        assert Path(args.csv_file).exists(), f"{args.csv_file} does not exist."
        metrics = pd.read_csv(args.csv_file).mean()
        sparsity = metrics['overall-sparsity']
        print(f"Average Sparsity: {sparsity:.3f}")
        for lb_key in ['50%-no-skip', '50%-skip', '25%-no-skip', '25%-skip']:
            load_balance = metrics[lb_key]
            total_lat = calc_sanger_latency(sparsity, load_balance, args.seq_len)
            print(f"Load Balance ({lb_key}): {load_balance:.3f}")
            print(f"Sanger Latency ({lb_key}): {total_lat * 1000:.3f} ms")


if __name__ == "__main__":
    main()