"""Shared helpers for the GQA 4-cases milestone benches.

Used by both ``milestone_gqa_decode_long_ctx_4cases`` and
``milestone_gqa_prefill_long_ctx_4cases``. Underscore-prefixed so the
benches package auto-discovery (which skips ``_*.py``) doesn't try to
audit it for a ``@bench`` decorator.
"""
from __future__ import annotations

from kernbench.ccl.install import load_ccl_config, resolve_algorithm_config


def _ccl_cfg() -> dict:
    """Resolve the lrab hierarchical AllReduce CCL config."""
    return resolve_algorithm_config(
        load_ccl_config(), name="lrab_hierarchical_allreduce",
    )


def _summarize_op_log(op_log) -> dict[str, int]:
    """Per-panel op_log counts (gemm, ipcq_copy, dma_read, dma_write)."""
    gemm_count = 0
    ipcq_copy_count = 0
    dma_read_count = 0
    dma_write_count = 0
    for r in op_log:
        if r.op_kind == "gemm":
            gemm_count += 1
        elif r.op_name == "dma_read":
            dma_read_count += 1
        elif r.op_name == "dma_write":
            dma_write_count += 1
        elif r.op_name == "ipcq_copy":
            ipcq_copy_count += 1
    return {
        "gemm_count": gemm_count,
        "ipcq_copy_count": ipcq_copy_count,
        "dma_read_count": dma_read_count,
        "dma_write_count": dma_write_count,
    }
