flatten_test.py

"""
Create tests for the flatten interface

Output format:
The .dat file has the following format (example with two outputs, dn n = 1,2)
<wen1> <wen2> <wdata> <wenexp1> <wenexp2> <wdataexp>
"""

import argparse
import random

from test_module import write_dat  # noqa: E402

random.seed(7357)

parser = argparse.ArgumentParser(description="Testing script for flatten")
parser.add_argument("-inp", help="Number of inputs", type=int)
parser.add_argument("-bw", help="Bit width", type=int)
parser.add_argument("-dw", help="Data width", type=int)
parser.add_argument("-ks", help="Kernel size", type=int)
args = parser.parse_args()

N_TESTS = 100

RANGE = int(2 ** (args.bw - 1) - 1)

# Create virtual counter
virt_cnt = [i % (args.dw + args.ks - 1) for i in range(N_TESTS)]

# Construct wen, addr, and data_i
wen = [[0 for _ in range(args.inp)] for _ in range(N_TESTS)]
waddr = [0 for _ in range(N_TESTS)]
data_i = [0 for _ in range(N_TESTS)]

start_list = [1] + [0 for _ in range(args.inp - 1)]

# Add wen and random data
ind = 0
for i in range(N_TESTS):
    addr = virt_cnt[i]
    if addr >= args.ks // 2 and addr <= args.dw + args.ks // 2 - 1:
        wen[i] = start_list[-ind:] + start_list[:-ind]
        data_i[i] = random.randint(-RANGE, RANGE)
        waddr[i] = addr - args.ks // 2

    if addr == args.dw - 1:
        ind += 1

wen_exp = [[0 for _ in range(args.inp * args.dw)] for _ in range(N_TESTS)]
waddr_exp = [0 for _ in range(N_TESTS)]
data_exp = data_i

count = 0

start_list = [1] + [0 for _ in range(args.dw * args.inp - 1)]
for i in range(N_TESTS):
    if not all(w == 0 for w in wen[i]):
        ind = 0
        wen_exp[i] = 1
        for j in range(args.inp):
            if wen[i][j] == 1:
                ind = j

        rot = waddr[i] + ind * args.dw
        wen_exp[i] = start_list[-rot:] + start_list[:-rot]

write_dat("flatten.dat", N_TESTS, data_i, wen, waddr, data_exp, wen_exp, waddr_exp)