import asyncio import math import os import cocotb from cocotb.triggers import RisingEdge, Timer import numpy as np MODE = os.getenv("TB_MODE", "int8").strip().lower() LANE_W = 8 DIM = 16 L = 2 async def gen_clk(sig, period_ns=10): while True: sig.value = 0 await Timer(period_ns // 2, unit="ns") sig.value = 1 await Timer(period_ns // 2, unit="ns") def to_signed(val, bits): if val & (1 << (bits - 1)): return val - (1 << bits) return val def float_to_fp8_e4m3(val): if val == 0: return 0 sign = 0 if val < 0: sign = 1 val = -val fp8_max = 448.0 val = min(val, fp8_max) exp_unbiased = math.floor(math.log2(val)) exp_biased = exp_unbiased + 7 if exp_biased <= 0: exp_biased = 0 mantissa = round(val / (2 ** (-6)) * 8) mantissa = min(mantissa, 7) else: exp_biased = min(exp_biased, 14) mantissa_f = val / (2 ** (exp_biased - 7)) - 1.0 mantissa = round(mantissa_f * 8) if mantissa >= 8: mantissa = 7 result = (sign << 7) | (exp_biased << 3) | (mantissa & 0x7) return result & 0xFF def fp8_e4m3_to_float(byte): byte = int(byte) & 0xFF sign = (byte >> 7) & 1 exp = (byte >> 3) & 0xF mant = byte & 0x7 if exp == 0b1111 and mant == 0b111: return float("nan") if exp == 0: val = (mant / 8.0) * (2.0 ** (-6)) else: val = (1.0 + mant / 8.0) * (2.0 ** (exp - 7)) return -val if sign else val def pack_row_dim(row, width, dim, m): v = 0 for i in range(m): x = row[i] if i < dim else 0 if MODE == "fp8": x_u = float_to_fp8_e4m3(x) else: x_u = int(x) & ((1 << width) - 1) v |= x_u << (i * width) return v def decode_value(val, bits): if MODE == "fp8": return fp8_e4m3_to_float(val) return to_signed(val, bits) def model_value(x): if MODE == "fp8": return fp8_e4m3_to_float(float_to_fp8_e4m3(x)) return to_signed(int(x) & 0xFF, 8) def pad_matrix(mat, size): out = [[0 for _ in range(size)] for _ in range(size)] for i, row in enumerate(mat[:size]): for j, val in enumerate(row[:size]): out[i][j] = val return out async def reset(dut): dut.i_rst_n.value = 0 dut.i_break.value = 0 dut.i_start.value = 0 dut.i_last.value = 0 dut.i_ab_valid.value = 0 dut.i_d_ready.value = 1 dut.i_a.value = 0 dut.i_b.value = 0 dut.i_dim.value = 0 for _ in range(5): await RisingEdge(dut.i_clk) dut.i_rst_n.value = 1 await RisingEdge(dut.i_clk) async def drive_hb(dut, list_a, list_b, runtime_dim, m, need_break, width_ab): for idx, (a_mat, b_mat) in enumerate(zip(list_a, list_b)): for i in range(runtime_dim): dut.i_a.value = pack_row_dim(a_mat[i], width_ab, runtime_dim, m) dut.i_b.value = pack_row_dim(b_mat[i], width_ab, runtime_dim, m) dut.i_start.value = 1 if idx == 0 and i == 0 else 0 dut.i_last.value = 1 if idx == len(list_a) - 1 and i == 0 else 0 dut.i_dim.value = runtime_dim dut.i_ab_valid.value = 1 dut.i_break.value = need_break and (idx == len(list_a) - 1) and (i == runtime_dim - 1) while True: await RisingEdge(dut.i_clk) if dut.o_ab_ready.value: break dut.i_ab_valid.value = 0 dut.i_break.value = 0 dut.i_start.value = 0 dut.i_last.value = 0 async def collect_outputs(dut, width_cd, m, runtime_dim, storage): current_rows = [] txn = 0 while True: await RisingEdge(dut.i_clk) if dut.o_d_valid.value and dut.i_d_ready.value: d_val = int(dut.o_d.value) row = [] for j in range(m): val_u = (d_val >> (j * width_cd)) & ((1 << width_cd) - 1) row.append(decode_value(val_u, width_cd)) current_rows.append(row) if len(current_rows) == runtime_dim: storage[txn] = current_rows.copy() current_rows.clear() txn += 1 def build_test_stream(): if MODE == "fp8": a = [[0.25 + 0.25 * (((7 * i + 3 * j) % 5) % 4) for j in range(16)] for i in range(16)] b = [[0.25 + 0.25 * (((5 * i + 11 * j + 1) % 5) % 4) for j in range(16)] for i in range(16)] else: a = [[((i * 2 + j) % 4) for j in range(16)] for i in range(16)] b = [[((i + 2 * j) % 4) for j in range(16)] for i in range(16)] return [(a, b)] def print_matrix(title, mat): print(f"{title}:") for row in mat: if MODE == "fp8": print(" ".join(f"{float(x):4.2f}" for x in row)) else: print(" ".join(f"{int(x):2d}" for x in row)) @cocotb.test() async def test_mmacu(dut): width_ab = LANE_W m = len(dut.i_a) // width_ab width_cd = len(dut.o_d) // m runtime_dim = min(DIM, m) cocotb.start_soon(gen_clk(dut.i_clk)) await reset(dut) outputs = {} cocotb.start_soon(collect_outputs(dut, width_cd, m, runtime_dim, outputs)) (a, b) = build_test_stream()[0] await drive_hb(dut, [a], [b], runtime_dim, m, need_break=True, width_ab=width_ab) for _ in range(4 * m + L): await RisingEdge(dut.i_clk) if 0 not in outputs: raise AssertionError("mmacu transaction 0 was not collected") a_sub = pad_matrix(a, m)[:runtime_dim] b_sub = pad_matrix(b, m)[:runtime_dim] if MODE == "fp8": a_np = np.array([[model_value(x) for x in row] for row in a_sub], dtype=np.float32) b_np = np.array([[model_value(x) for x in row] for row in b_sub], dtype=np.float32) golden = np.array((a_np.T @ b_np).astype(np.float32), dtype=np.float32)[:runtime_dim, :runtime_dim] hw = np.array(outputs[0], dtype=np.float32)[:runtime_dim, :runtime_dim] passed = np.allclose(golden, hw, rtol=0, atol=2) else: a_np = np.array([[model_value(x) for x in row] for row in a_sub], dtype=np.int32) b_np = np.array([[model_value(x) for x in row] for row in b_sub], dtype=np.int32) golden = (a_np.T @ b_np).astype(np.int32) golden = np.vectorize(lambda x: to_signed(x & 0xFF, 8))(golden) golden = np.array(golden, dtype=np.int32)[:runtime_dim, :runtime_dim] hw = np.array(outputs[0], dtype=np.int32)[:runtime_dim, :runtime_dim] passed = np.array_equal(golden, hw) print_matrix("Input A", a_sub) print_matrix("Input B", b_sub) print_matrix("Hardware Output", hw) print("\n" + "=" * 80) print(f"Verification against NumPy Model ({MODE.upper()})") print("=" * 80) print("Transaction 0: [PASSED]" if passed else "Transaction 0: [FAILED]") if not passed: print("Golden (Expected):") print(golden) print("Hardware (Actual):") print(hw) raise AssertionError("mmacu verification failed") print("\nALL TRANSACTIONS VERIFIED SUCCESSFULLY")