Source code for arcade.hardware.estimator

"""Analytical resource estimation per classifier type — data-driven approach."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import numpy as np

from .._logging import get_logger
from .base import BaseHardwareEstimator, ResourceReport, register_estimator  # noqa: F401 (re-export)

log = get_logger(__name__)


# ---------------------------------------------------------------------------
# Estimator parameter table
# ---------------------------------------------------------------------------
# Each entry maps a classifier name to:
#   extract_fn:  callable(model) -> (units: int, extra: dict)
#   lut_mul:     LUT = units * lut_mul
#   ff_mul:      FF  = units * ff_mul
#   dsp_mul:     DSP = units * dsp_mul
#   bram_div:    BRAM = max(1, units // bram_div)  (0 => bram=0)
#   base_latency:  base additive latency cycles
#   cycles_per_unit: latency_cycles = base_latency + units * cycles_per_unit
#                    (float; 0 means latency comes from extra)


@dataclass(frozen=True)
class _EstParams:
    lut_mul: float
    ff_mul: float
    dsp_mul: float
    bram_div: int  # 0 means bram=0
    base_latency: int
    cycles_per_unit: float


def _extract_nn(model: Any) -> tuple[int, dict]:
    import torch.nn as nn

    total_macs = 0
    total_weights = 0
    n_layers = 0
    for module in model.modules():
        if isinstance(module, nn.Linear):
            macs = module.in_features * module.out_features
            total_macs += macs
            total_weights += macs + module.out_features
            n_layers += 1
    return total_macs, {"weights": total_weights, "n_layers": n_layers, "macs": total_macs}


def _extract_threshold(model: Any) -> tuple[int, dict]:
    fs = getattr(model, "filter_set", None)
    n_filters = fs.total_features if fs else 1
    n_qubits = fs.num_qubits if fs else 1
    return n_filters, {"n_qubits": n_qubits}


def _extract_sklearn(model: Any) -> tuple[int, dict]:
    inner = getattr(model, "_model", None)
    if inner is None:
        return 0, {}
    n_params = 0
    if hasattr(inner, "coef_"):
        n_params = int(np.prod(inner.coef_.shape))
    elif hasattr(inner, "support_vectors_"):
        n_params = int(np.prod(inner.support_vectors_.shape))
    elif hasattr(inner, "coefs_"):
        n_params = int(sum(np.prod(c.shape) for c in inner.coefs_))
    return n_params, {}


def _extract_gmm(model: Any) -> tuple[int, dict]:
    n_classes = len(getattr(model, "_classes", []))
    n_components = getattr(model, "n_components", 1)
    models = getattr(model, "_models", {})
    total_params = 0
    for gmm in models.values():
        if hasattr(gmm, "means_"):
            d = gmm.means_.shape[1]
            total_params += n_components * (d + d * d)
    return total_params, {"n_classes": n_classes}


def _extract_ngrc(model: Any) -> tuple[int, dict]:
    w = getattr(model, "window_size", 10)
    deg = getattr(model, "polynomial_degree", 2)
    n_feat = max(100, w * 4 * (deg**2))
    ro = getattr(model, "_readout", None)
    if ro is not None and "W_out" in ro:
        n_feat = int(np.prod(ro["W_out"].shape))
    elif getattr(model, "_per_qubit_readouts", None):
        n_feat = sum(int(np.prod(r["W_out"].shape)) for r in model._per_qubit_readouts)
    return n_feat, {}


def _extract_path_signature(model: Any) -> tuple[int, dict]:
    depth = getattr(model, "signature_depth", 5)
    return 4**depth, {}


def _extract_transformer(model: Any) -> tuple[int, dict]:
    inner = getattr(model, "_model", model)
    d_model = getattr(model, "d_model", None) or getattr(inner, "d_model", 64)
    nhead = getattr(model, "nhead", None) or getattr(inner, "nhead", 4)
    num_layers = getattr(model, "num_layers", None) or getattr(inner, "num_layers", 2)
    seq_len = getattr(model, "seq_len", None) or getattr(inner, "seq_len", 16)

    attn_macs = num_layers * (4 * seq_len * d_model**2 + seq_len**2 * d_model)
    ffn_macs = num_layers * 2 * seq_len * d_model * 4 * d_model
    total_macs = attn_macs + ffn_macs
    total_weights = num_layers * (4 * d_model**2 + 2 * d_model * 4 * d_model)
    latency = num_layers * seq_len + seq_len

    log.info(
        "Transformer estimate: d_model=%d, nhead=%d, layers=%d, seq=%d, MACs=%d",
        d_model, nhead, num_layers, seq_len, total_macs,
    )
    return total_macs, {"weights": total_weights, "latency": latency, "macs": total_macs,
                        "note": "Analytical estimate; actual synthesis may vary significantly for transformers."}


def _extract_multilevel(model: Any) -> tuple[int, dict]:
    """Estimate MACs for leakage / multilevel (3-level) neural heads."""
    inner = getattr(model, "_model", model)
    total_macs = 0
    total_weights = 0
    n_layers = 0
    try:
        import torch.nn as nn
        for module in inner.modules():
            if isinstance(module, nn.Linear):
                macs = module.in_features * module.out_features
                total_macs += macs
                total_weights += macs + module.out_features
                n_layers += 1
    except ImportError:
        total_macs, total_weights, n_layers = 1000, 1000, 3

    if total_macs == 0:
        total_macs, total_weights, n_layers = 1000, 1000, 3

    qf = 1.5
    return total_macs, {"weights": total_weights, "n_layers": n_layers,
                        "macs": total_macs, "qutrit_factor": qf,
                        "note": "Qutrit (3-level) system; precision requirements higher than qubit."}


# ---------------------------------------------------------------------------
# Parameter table: maps classifier name -> (extract_fn, params)
# For types whose resource formulas don't fit the simple linear model,
# _build_report applies type-specific overrides using the extra dict.
# ---------------------------------------------------------------------------

ESTIMATOR_PARAMS: dict[str, tuple[Any, _EstParams]] = {
    "nn":             (_extract_nn,             _EstParams(lut_mul=3,  ff_mul=2,  dsp_mul=1,    bram_div=1024, base_latency=1, cycles_per_unit=0)),
    "threshold":      (_extract_threshold,      _EstParams(lut_mul=50, ff_mul=16, dsp_mul=1,    bram_div=0,    base_latency=2, cycles_per_unit=0)),
    "svm":            (_extract_sklearn,         _EstParams(lut_mul=5,  ff_mul=2,  dsp_mul=0.1,  bram_div=512,  base_latency=3, cycles_per_unit=0.01)),
    "qda":            (_extract_sklearn,         _EstParams(lut_mul=5,  ff_mul=2,  dsp_mul=0.1,  bram_div=512,  base_latency=3, cycles_per_unit=0.01)),
    "lda":            (_extract_sklearn,         _EstParams(lut_mul=5,  ff_mul=2,  dsp_mul=0.1,  bram_div=512,  base_latency=3, cycles_per_unit=0.01)),
    "gnb":            (_extract_sklearn,         _EstParams(lut_mul=5,  ff_mul=2,  dsp_mul=0.1,  bram_div=512,  base_latency=3, cycles_per_unit=0.01)),
    "mlp":            (_extract_sklearn,         _EstParams(lut_mul=5,  ff_mul=2,  dsp_mul=0.1,  bram_div=512,  base_latency=3, cycles_per_unit=0.01)),
    "gmm":            (_extract_gmm,            _EstParams(lut_mul=8,  ff_mul=3,  dsp_mul=0.2,  bram_div=256,  base_latency=0, cycles_per_unit=0)),
    # "reservoir" remains as deprecated alias of ngrc for one release.
    "reservoir":      (_extract_ngrc,            _EstParams(lut_mul=2,  ff_mul=1,  dsp_mul=0.25, bram_div=0,    base_latency=0, cycles_per_unit=0.1)),
    "ngrc":           (_extract_ngrc,            _EstParams(lut_mul=2,  ff_mul=1,  dsp_mul=0.25, bram_div=0,    base_latency=0, cycles_per_unit=0.1)),
    "path_signature": (_extract_path_signature,  _EstParams(lut_mul=10, ff_mul=5,  dsp_mul=0,    bram_div=0,    base_latency=0, cycles_per_unit=1)),
    "transformer":    (_extract_transformer,     _EstParams(lut_mul=2,  ff_mul=0,  dsp_mul=0.1,  bram_div=0,    base_latency=0, cycles_per_unit=0)),
    "leakage":        (_extract_multilevel,      _EstParams(lut_mul=3,  ff_mul=2,  dsp_mul=1,    bram_div=1024, base_latency=1, cycles_per_unit=0)),
    "multilevel":     (_extract_multilevel,      _EstParams(lut_mul=3,  ff_mul=2,  dsp_mul=1,    bram_div=1024, base_latency=1, cycles_per_unit=0)),
}


def _build_report(
    name: str,
    units: int,
    params: _EstParams,
    extra: dict,
    clock_ns: float,
) -> ResourceReport:
    """Construct a ResourceReport from extracted units, params, and extras."""
    qf = extra.get("qutrit_factor", 1.0)
    weights = extra.get("weights", units)
    n_layers = extra.get("n_layers", 0)

    lut = int(units * params.lut_mul * qf)
    ff = int(units * params.ff_mul * qf) if params.ff_mul else int(weights * 2 * qf) if weights else 0
    dsp = max(1, int(units * params.dsp_mul * qf)) if params.dsp_mul else 0
    bram = max(1, int(weights * qf) // params.bram_div) if params.bram_div else extra.get("bram", 0)

    # Threshold special: add qubit-dependent LUT
    if "n_qubits" in extra:
        lut += extra["n_qubits"] * 20

    # Latency: use extra override, or formula, or layer-based
    if "latency" in extra:
        latency_cycles = extra["latency"]
    elif params.cycles_per_unit:
        latency_cycles = max(params.base_latency, int(params.base_latency + units * params.cycles_per_unit))
    elif n_layers:
        latency_cycles = n_layers + params.base_latency
    elif name == "gmm":
        latency_cycles = extra.get("n_classes", 1) * 5
    elif name in ("reservoir", "ngrc"):
        latency_cycles = max(1, units // 10)
    elif name == "path_signature":
        latency_cycles = units
    else:
        latency_cycles = params.base_latency or 1

    # BRAM for ngrc/path_signature (reservoir shares ngrc costing)
    if name in ("ngrc", "reservoir"):
        bram = 1
    if name == "path_signature":
        bram = 2

    macs = extra.get("macs", None)
    note = extra.get("note", "")

    if macs and name in ("nn", "leakage", "multilevel", "transformer"):
        log.info("%s estimate (%s): %d MACs, DSP=%d, LUT=%d", name, name, macs, dsp, lut)

    return ResourceReport(
        classifier_name=name,
        lut=lut,
        ff=ff,
        bram=bram,
        dsp=dsp,
        latency_cycles=latency_cycles,
        latency_ns=latency_cycles * clock_ns,
        macs_per_inference=int(macs) if macs else None,
        clock_period_ns=clock_ns,
        estimated=True,
        note=note,
    )


# ---------------------------------------------------------------------------
# Public utilities
# ---------------------------------------------------------------------------



def _estimate_nn(model: Any, *, clock_ns: float = 5.0, name: str = "nn") -> ResourceReport:
    """Estimate resources for a ``torch.nn.Module`` (used by unit tests)."""
    extract_fn, params = ESTIMATOR_PARAMS["nn"]
    units, extra = extract_fn(model)
    return _build_report(name, units, params, extra, clock_ns)


def count_linear_macs(model: Any) -> int | None:
    """Naïvely sum dense MAC counts over ``torch.nn.Linear`` layers."""
    try:
        import torch.nn as nn
    except ImportError:
        return None
    total = 0
    for module in model.modules():
        if isinstance(module, nn.Linear):
            total += module.in_features * module.out_features
    return int(total) if total > 0 else None


def _parse_precision_bits(precision: str) -> int:
    import re
    m = re.match(r"ap_fixed<(\d+)", precision)
    return int(m.group(1)) if m else 16


def _apply_hw_config_scaling(report: ResourceReport, config: dict) -> ResourceReport:
    """Adjust analytical estimates for reuse factor, precision, and I/O mode."""
    reuse = max(1, int(config.get("reuse_factor", 1)))
    bits = _parse_precision_bits(str(config.get("precision", "ap_fixed<16,6>")))
    bit_scale = bits / 16.0
    io_type = config.get("io_type", "io_parallel")
    io_area_scale = 1.15 if io_type == "io_stream" else 1.0
    area_scale = (1.0 / reuse) * bit_scale * io_area_scale
    lat_scale = float(reuse) * (1.0 / max(bit_scale, 0.25))

    def _scale_int(v: int | None) -> int | None:
        return int(v * area_scale) if v is not None else None

    cycles = int(report.latency_cycles * lat_scale) if report.latency_cycles is not None else None
    lat_ns = float(report.latency_ns * lat_scale) if report.latency_ns is not None else None

    return ResourceReport(
        classifier_name=report.classifier_name,
        lut=_scale_int(report.lut),
        ff=_scale_int(report.ff),
        bram=_scale_int(report.bram),
        dsp=_scale_int(report.dsp),
        latency_cycles=cycles,
        latency_ns=lat_ns,
        macs_per_inference=report.macs_per_inference,
        clock_period_ns=report.clock_period_ns,
        estimated=report.estimated,
        note=report.note,
    )


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------

[docs] def estimate_resources( model: Any, config: dict | None = None, *, classifier_name: str = "unknown", ) -> ResourceReport: """Estimate FPGA resources for a trained classifier. Dispatches to type-specific parameter sets based on the model type. Args: model: Trained classifier (NN, hybrid, threshold, sklearn, etc.). config: The ``hardware`` section of the ARCADE config (optional). classifier_name: Human-readable name for the report. Returns: A ``ResourceReport`` with resource estimates. """ if config is None: config = {} clock_ns = config.get("clock_period_ns", 5.0) # Try torch.nn.Module directly try: import torch.nn as nn if isinstance(model, nn.Module): extract_fn, params = ESTIMATOR_PARAMS["nn"] units, extra = extract_fn(model) report = _build_report(classifier_name, units, params, extra, clock_ns) return _apply_hw_config_scaling(report, config) except ImportError: pass # Try BaseClassifier wrapper from ..classifiers.base import BaseClassifier if isinstance(model, BaseClassifier): head = getattr(model, "_model", None) if head is not None: try: import torch.nn as nn if isinstance(head, nn.Module): nm = getattr(model, "name", classifier_name) extract_fn, params = ESTIMATOR_PARAMS["nn"] units, extra = extract_fn(head) report = _build_report(nm, units, params, extra, clock_ns) return _apply_hw_config_scaling(report, config) except ImportError: pass name = getattr(model, "name", classifier_name) if name in ESTIMATOR_PARAMS: extract_fn, params = ESTIMATOR_PARAMS[name] units, extra = extract_fn(model) if units == 0 and not extra: return ResourceReport(classifier_name=name, clock_period_ns=clock_ns) report = _build_report(name, units, params, extra, clock_ns) return _apply_hw_config_scaling(report, config) log.warning("No estimator for classifier %r; returning empty report", name) return ResourceReport(classifier_name=classifier_name, clock_period_ns=clock_ns) log.warning("No estimator for model type %s; returning empty report", type(model).__name__) return ResourceReport(classifier_name=classifier_name, clock_period_ns=clock_ns)
# --------------------------------------------------------------------------- # Registered estimator class # ---------------------------------------------------------------------------
[docs] @register_estimator("analytical") class AnalyticalEstimator(BaseHardwareEstimator): """Data-driven analytical hardware estimator."""
[docs] def estimate(self, model: Any, **kwargs: Any) -> ResourceReport: return estimate_resources(model, **kwargs)