Mini-Nav/hw/sim/tests/top/utils.py

# -*- coding: utf-8 -*-
"""
CAM 顶层测试的共享辅助函数。

被 no_noise/、write_noise/、read_noise/ 等配置目录的测试文件共同引用。
提供：
  - Verilator 参数读取（带回退推断）
  - DUT 复位、空闲等待
  - 行写入 / 批量写入（握手协议）
  - 查询发起 / Top-K 结果流收集
  - score_debug 解包（SIM_DEBUG 模式）

所有函数都是 async，调用者需要处于 cocotb 协程上下文中。
"""
from __future__ import annotations

import numpy as np
from cocotb.triggers import RisingEdge
from model.ref_model import (  # noqa: E402
    match_top1,
    unpack_score_debug_flat,
)

# ── 默认拓扑参数（当 Verilator 参数不可用时使用） ───────────────────────────
DEFAULT_NUM_ROWS = 4096
DEFAULT_HASH_BITS = 512
DEFAULT_LANES = 8
DEFAULT_SCORE_BITS = 10


# ═══════════════════════════════════════════════════════════════════════════════
# 参数读取（从 Verilator 参数或信号宽度推断）
# ═══════════════════════════════════════════════════════════════════════════════


def get_param(dut, name, default=None):
    """从 DUT 读取 Verilator 暴露的参数值，失败则返回 default。

    某些参数仅在特定编译配置下暴露（如 SIM_DEBUG 下的 score_debug_flat），
    此时回退到 default 是预期行为而非错误。
    """
    try:
        val = getattr(dut, name, None)
        if val is not None:
            return int(val.value)
    except Exception:
        pass
    return default


def dut_num_rows(dut):
    """获取 NUM_ROWS：优先读参数，否则从 wr_addr 位宽推断 (NUM_ROWS = 2^ROW_BITS)。"""
    val = get_param(dut, "NUM_ROWS", None)
    if val is not None:
        return val
    return 1 << len(dut.wr_addr)


def dut_hash_bits(dut):
    """获取 HASH_BITS：优先读参数，否则从 write_hash 信号位宽推断。"""
    val = get_param(dut, "HASH_BITS", None)
    if val is not None:
        return val
    return len(dut.write_hash)


def dut_lanes(dut):
    """获取 LANES：优先读参数，否则从 rd_resp_row_ids / wr_addr 位宽推断。"""
    val = get_param(dut, "LANES", None)
    if val is not None:
        return val
    return len(dut.rd_resp_row_ids) // len(dut.wr_addr)


def dut_score_bits(dut):
    """获取 SCORE_BITS：优先读参数，否则从 top1_score 信号位宽推断。"""
    val = get_param(dut, "SCORE_BITS", None)
    if val is not None:
        return val
    return len(dut.top1_score)


# ═══════════════════════════════════════════════════════════════════════════════
# 协议层辅助函数（复位 / 空闲等待 / 行写入 / 查询）
# ═══════════════════════════════════════════════════════════════════════════════


async def reset_dut(dut):
    """复位 DUT：rst_n 拉低 5 周期，释放后再等 2 周期。

    所有控制信号在复位期间保持无效电平。
    result_ready 初始化为 1（准备接收结果）。
    """
    dut.rst_n.value = 0
    dut.wr_valid.value = 0
    dut.wr_addr.value = 0
    dut.write_hash.value = 0
    dut.query_valid.value = 0
    dut.query_hash.value = 0
    dut.result_ready.value = 1

    for _ in range(5):
        await RisingEdge(dut.clk)

    dut.rst_n.value = 1

    for _ in range(2):
        await RisingEdge(dut.clk)


async def wait_idle(dut):
    """等待 DUT 完全空闲：wr_ready=1 且 query_ready=1。

    这是发起新操作前的前置条件——CAM 是半双工的，
    同一时刻只能进行写入或查询。
    """
    while not (int(dut.wr_ready.value) and int(dut.query_ready.value)):
        await RisingEdge(dut.clk)


async def write_row(dut, addr, value):
    """写入单行：使用 wr_valid/wr_ready 握手协议。

    流程：等待空闲 → 驱动地址和数据 → 等待握手完成 → 等待写流水线排空。
    """
    await wait_idle(dut)

    dut.wr_addr.value = addr
    dut.write_hash.value = int(value)
    dut.wr_valid.value = 1

    while True:
        await RisingEdge(dut.clk)
        if int(dut.wr_ready.value):
            break

    dut.wr_valid.value = 0
    await wait_idle(dut)


async def write_rows(dut, rows):
    """按顺序写入所有行（行索引 = 数组下标）。"""
    for idx, value in enumerate(rows):
        await write_row(dut, idx, value)


async def collect_topk(dut, timeout_cycles: int = 2000):
    """收集串行 Top-K 结果流的所有 beat。

    保持 result_ready=1，逐个时钟周期采样 result_valid，
    直到 result_last 被断言。

    返回：[(rank, row, score, last), ...]  列表
    超时则抛出 AssertionError。

    注意：此函数会「消耗」整个结果流，调用后 DUT 回到空闲状态。
    """
    dut.result_ready.value = 1
    beats = []
    for _ in range(timeout_cycles):
        if int(dut.result_valid.value):
            rank = int(dut.result_rank.value)
            row = int(dut.result_row.value)
            score = int(dut.result_score.value)
            last = int(dut.result_last.value)
            beats.append((rank, row, score, last))
            if last:
                return beats
        await RisingEdge(dut.clk)
    raise AssertionError("Top-K result stream did not finish")


async def query_topk_once(dut, query, timeout_cycles=2000):
    """发起一次查询并收集完整的串行 Top-K 结果流。

    完整流程：
      1. 等待 DUT 空闲
      2. 通过 query_valid/query_ready 握手发送查询
      3. 消费完整的结果流
      4. 读取 score_debug_flat（如果存在）

    返回：(beats, top1_index, top1_score, score_debug)
      - beats: [(rank, row, score, last), ...]
      - score_debug: np.ndarray 或 None（SIM_DEBUG 模式）
    """
    await wait_idle(dut)

    dut.query_hash.value = int(query)
    dut.query_valid.value = 1

    # 等待查询握手完成
    while True:
        await RisingEdge(dut.clk)
        if int(dut.query_ready.value):
            break

    dut.query_valid.value = 0

    # 消费完整串行结果流
    beats = await collect_topk(dut, timeout_cycles=timeout_cycles)

    # score_debug 在查询完成后可用（需 SIM_DEBUG 编译）
    num_rows = dut_num_rows(dut)
    score_bits = dut_score_bits(dut)
    score_debug = None
    if hasattr(dut, "score_debug_flat"):
        score_debug = unpack_score_debug_flat(
            int(dut.score_debug_flat.value),
            num_rows,
            score_bits,
        )

    return beats, beats[0][1], beats[0][2], score_debug


async def query_once(dut, query):
    """发起查询，返回 (top1_index, top1_score, score_debug)。

    内部调用 query_topk_once 并消费完整结果流，仅保留 rank-0 数据。
    """
    _, top1_index, top1_score, score_debug = await query_topk_once(dut, query)
    return top1_index, top1_score, score_debug


# ═══════════════════════════════════════════════════════════════════════════════
# 便捷验证函数
# ═══════════════════════════════════════════════════════════════════════════════


def assert_baseline_top1(query, rows, top1_index, top1_score, hash_bits,
                         query_index, score_debug=None):
    """验证无噪声场景下的基线 Top-1 结果。

    检查项：
      1. Top-1 与 match_top1 参考模型一致
      2. 查询自身所在行 → score == hash_bits（完全匹配）
      3. score_debug 数组（如果存在）与模型一致
    """
    expected = match_top1(query, rows, width=hash_bits)
    assert top1_index == expected.top1_index, (
        f"top1_index mismatch: {top1_index} != {expected.top1_index}"
    )
    assert top1_score == expected.top1_score, (
        f"top1_score mismatch: {top1_score} != {expected.top1_score}"
    )
    assert top1_index == query_index, (
        f"Expected query_index={query_index} to match self, got top1_index={top1_index}"
    )
    assert top1_score == hash_bits, (
        f"Self-query should score {hash_bits}, got {top1_score}"
    )

    if score_debug is not None:
        assert np.array_equal(score_debug, expected.scores), (
            "score_debug does not match model scores"
        )