"""Ascend950PR (DAV_3510) 硬件规格参数.

数值来源: 《BMM算子优化分析 v0.98》§二 + 昇腾950 NPU架构白皮书.
换芯片时逻辑结构不变, 只需新增一份同结构参数表.

单位约定: 算力 FLOP/s, 带宽 Byte/s, 容量 Byte, 时延 秒.

dtype 感知算力 (issue#28, 设计文档 docs/05 §2):
  Cube 算力按输入 dtype 分档, 基准 = BF16 (fp16 同速); 白皮书: FP8/MXFP8/HiF8
  提供 2x FP16 张量 TFLOPS, MXFP4 提供 4x FP16; FP32/TF32 同代比值为假设值
  (按 DaVinci 惯例 = 1/2, 白皮书未给同代比值), int8 假设同 FP8, 均待实测标定.
  AIV 逐元素通量按 lane 位宽等比假设 (16bit x2 / 8bit x4 / 4bit x8, 待标定).
"""

from __future__ import annotations

from dataclasses import dataclass

# Cube 精度因子 (相对 BF16/FP16 基准档; A/B 不一致时取较慢一侧 = min 因子)
CUBE_DTYPE_FACTOR = {
    "fp32": 0.5, "f32": 0.5, "tf32": 0.5,   # 假设 = 1/2, 待实测标定
    "fp16": 1.0, "f16": 1.0, "bf16": 1.0,
    "fp8": 2.0, "fp8_e4m3": 2.0, "fp8_e5m2": 2.0, "int8": 2.0,  # 白皮书 FP8=2xFP16; int8 假设同 FP8
    "fp4": 4.0, "fp4_e2m1": 4.0,             # 白皮书 MXFP4=4xFP16; 普通 fp4 假设同 MXFP4
}

# AIV 逐元素通量因子 (相对 fp32 128 lane/拍/核; 位宽等比假设, 待实测标定)
AIV_DTYPE_FACTOR = {
    "fp32": 1.0, "f32": 1.0, "tf32": 1.0,
    "fp16": 2.0, "f16": 2.0, "bf16": 2.0,
    "fp8": 4.0, "fp8_e4m3": 4.0, "fp8_e5m2": 4.0, "int8": 4.0,
    "fp4": 8.0, "fp4_e2m1": 8.0,
}

_RATE_FALLBACK = 1.0  # 未知 dtype 按基准档 (models.dtype_bytes 已先行校验, 正常不会到达)


@dataclass(frozen=True)
class NpuSpec:
    name: str = "Ascend950PR"

    # ---- 核数与算力 ----
    aic_num: int = 32                      # C: AIC (Cube) 核数
    aiv_num: int = 64                      # AIV (Vector) 核数
    cube_peak_tflops: float = 486.0        # 全芯片 Cube BF16 标称算力 (乘加各计一次)
    aiv_freq_ghz: float = 1.65
    aiv_fp32_per_cycle: int = 128          # 单 AIV 每拍 fp32 求和吞吐

    # ---- 片上存储 (每核) ----
    l1_bytes: int = 512 * 1024             # L1 Buffer 512KB
    l0a_bytes: int = 64 * 1024             # L0A 64KB
    l0b_bytes: int = 64 * 1024             # L0B 64KB
    l0c_bytes: int = 256 * 1024            # L0C 256KB (fp32 累加, 4B/元素)

    # ---- L2 / GM ----
    l2_bytes: int = 128 * 1024 * 1024      # L2 Cache 128MB
    bw_l2: float = 5.2e12                  # L2 读写带宽 5.2TB/s
    bw_gm: float = 1.6e12                  # GM(HBM) 带宽 1.6TB/s (读写共享)

    # ---- 搬移效率经验约束 (950PR 实测总结) ----
    dvalue_recommend: int = 512            # dValue 推荐值 (Byte), 更大无额外收益
    dvalue_min: int = 128                  # dValue 下限 (Byte)
    dvalue_hw_min: int = 256               # DMA 硬件突发下限 (Byte) —— 尾轮文档 §2.3
    min_tile_size: int = 16 * 1024         # min_TileSize: 单块搬移最小量 16KB
    min_datamount_per_core: int = 480 * 1024   # min_DatamountPerCore: 单核搬移总量下限 480KB
    min_core_num_ratio: float = 0.8        # minCoreNum ≈ 0.8 * C (经验: 约 3/4 核并发才达 90%+ 带宽利用率)

    # ---- DMA 固定开销 ----
    t_cmd_ns: float = 0.0                  # T_cmd: 单次 GM->L1 DMA 命令固定开销 (ns)。
                                           # 未标定, 按 0 处理 (issue#36 用户澄清):
                                           # MergeBatch 的合并收益由搬移效率模型
                                           # (move_eff: 单块 tile=nValue*dValue*dt
                                           # 放大 b0 倍) 刻画, 不依赖 T_cmd 估计值.

    # ---- 带宽模型假设 (issue#26) ----
    # GM 1.6TB/s 为读写共享总线 (读+写累加计时); L2 带宽读写各自独享 5.2TB/s;
    # active_cores < C 时按"每核份额线性配平"为模型假设 (如 1 核也给 50GB/s),
    # 真实低核数带宽利用率低于线性, 需 msProf 实测曲线标定后替换该假设.

    # ---- Cube 计算粒度 ----
    fractal: int = 16                      # 16x16x16 基本块

    # ---- 派生量 (属性) ----
    @property
    def q16(self) -> float:
        """单核 Cube BF16/FP16 峰值算力 (FLOP/s) = 15.1875T."""
        return self.cube_peak_tflops * 1e12 / self.aic_num

    @staticmethod
    def _dtype_key(dtype) -> str:
        return str(dtype).strip().lower()

    def cube_factor(self, dtype_a, dtype_b=None) -> float:
        """按输入 dtype 取 Cube 精度因子 (issue#28).

        A/B 不一致时取**较慢一侧** (因子较小者, 等价字节较大者);
        Cube 乘法两侧的实际吞吐受较慢精度限制.
        """
        keys = [self._dtype_key(dtype_a)]
        if dtype_b is not None:
            keys.append(self._dtype_key(dtype_b))
        factors = [CUBE_DTYPE_FACTOR.get(k, _RATE_FALLBACK) for k in keys]
        return min(factors)

    def q_cube(self, dtype_a, dtype_b=None) -> float:
        """单核 Cube 峰值算力 (FLOP/s), 按输入 dtype 分档 (issue#28)."""
        return self.q16 * self.cube_factor(dtype_a, dtype_b)

    def aiv_elem_factor(self, dtype) -> float:
        """按 dtype 取 AIV 逐元素通量因子 (相对 fp32 基准)."""
        return AIV_DTYPE_FACTOR.get(self._dtype_key(dtype), _RATE_FALLBACK)

    @property
    def aiv_elem_rate_fp32(self) -> float:
        """AIV fp32 逐元素吞吐 (元素/s), 64 核合计."""
        return self.aiv_num * self.aiv_fp32_per_cycle * self.aiv_freq_ghz * 1e9

    def aiv_elem_rate(self, dtype) -> float:
        """AIV 逐元素吞吐 (元素/s, 64 核合计), 按 dtype 分档 (issue#28)."""
        return self.aiv_elem_rate_fp32 * self.aiv_elem_factor(dtype)

    @property
    def q_aiv(self) -> float:
        """AIV fp32 逐元素吞吐 (元素/s, 64 核合计) = aiv_elem_rate_fp32.

        注意: 该值只适用于 fp32 数据 (如 StreamK 对 fp32 部分和求和);
        逐元素运算按输入 dtype 用 aiv_elem_rate(dtype) (issue#28).
        """
        return self.aiv_elem_rate_fp32

    @property
    def bw_pc(self) -> float:
        """单核 GM 带宽份额 (Byte/s) = W_GM / C = 50GB/s."""
        return self.bw_gm / self.aic_num

    @property
    def bw_l2_pc(self) -> float:
        """单核 L2 带宽份额 (Byte/s) = W_L2 / C ≈ 162.5GB/s."""
        return self.bw_l2 / self.aic_num

    @property
    def r16(self) -> float:
        """16bit 位宽平衡点算存比 R_16 = Cube峰值 / (GM带宽/2B) ≈ 607.5 FLOP/元素.

        按 issue#28 速率表该比值对全 dtype 不变 (Cube 因子与元素字节数互成反比),
        故入口条件沿用单一 R_16.
        """
        return self.cube_peak_tflops * 1e12 / (self.bw_gm / 2)

    @property
    def min_core_num(self) -> int:
        """确保高带宽利用率的并行搬移核数下限."""
        return round(self.min_core_num_ratio * self.aic_num)

    @property
    def t_cmd(self) -> float:
        """T_cmd (秒)."""
        return self.t_cmd_ns * 1e-9

    @property
    def l0c_elems(self) -> int:
        """L0C 容量 (fp32 元素数) = 65536."""
        return self.l0c_bytes // 4


ASCEND950PR = NpuSpec()