"""
=============================================================================
CS349E · 前置讲 P5 · 01_links_and_topologies.py —— 链路的 α–β 模型与网络拓扑
=============================================================================

【这个脚本要回答的问题】
    芯片之间的"网"由两样东西决定：每条**链路**有多快（延迟 α、带宽 β），以及链路怎样连成**拓扑**
    （交换机、环、torus、两级网络）。本脚本：

      演示 1：α–β 模型——发一条消息要 α + n/β。几种常见链路（NVLink、PCIe、InfiniBand、TPU ICI、以太网），
              消息从 1 KB 到 1 GB 时的耗时和"实际拿到的带宽"；半带宽点 n½ = α·β。
      演示 2：拓扑——交换机、双向环、二维 torus、两级网络：每个设备几个端口、最远要几跳（直径）、对分带宽。
      演示 3：点对点发送 1 GB 和 4 KB：同一个节点里、跨节点、跨 PCIe……差多少？

【怎么读输出】全部由公式和模拟器算出，每次运行结果相同。链路的带宽是公开规格（单向）；延迟 α 是量级上的假设值。

【运行】python3 01_links_and_topologies.py      （约 1 秒）
【依赖】numpy；从 ../lecture01 导入 console、hw_specs；同目录的 comm_sim。
"""

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from console import fmt_bytes, fmt_time, print_header, print_observations, print_table  # noqa: E402
from hw_specs import H100, TPU_V5E  # noqa: E402
from comm_sim import (ETH100, IB_NDR, ICI_V5E, LINKS, NVLINK4, PCIE5, Comm, bisection_bw, ring, switch,  # noqa: E402
                      torus2d, two_level)


# ---------------------------------------------------------------------------
# 演示 1：α–β 模型
# ---------------------------------------------------------------------------
def demo1_alpha_beta():
    """
    T(n) = α + n/β；实际拿到的带宽 = n / T(n)。消息小于 n½ = α·β 时，拿到的带宽不到一半。

    【数字从哪来】带宽：NVLink 4 每块 H100 18 条 × 25 GB/s = 450 GB/s 单向（和第 1 讲 hw_specs 的 H100.link_bw 一致）；
        InfiniBand NDR 400 Gb/s = 50 GB/s；TPU v5e 每条 ICI 链路 45 GB/s 单向（hw_specs 里 4 条合计）；PCIe 5.0 x16 约 63 GB/s。
        延迟：NVLink / PCIe 2 µs、InfiniBand 5 µs、ICI 1 µs（How to Scale Your Model 的量级）、以太网 10 µs——都是假设值。
    """
    print_header("演示 1：α–β 模型——T(n) = α + n/β，小消息吃延迟，大消息吃带宽")
    assert NVLINK4.bw == H100.link_bw and abs(4 * ICI_V5E.bw - TPU_V5E.link_bw) < 1
    sizes = [1e3, 64e3, 1e6, 64e6, 1e9]
    rows = []
    for L in LINKS:
        cells = [L.name, f"{L.bw / 1e9:.1f} GB/s", f"{L.alpha * 1e6:.0f} µs", fmt_bytes(L.half_bw_size())]
        for n in sizes:
            cells.append(f"{n / L.time(n) / 1e9:.1f}")
        rows.append(cells)
    print_table(["链路", "带宽 β（单向）", "延迟 α（假设）", "半带宽点 n½ = αβ"] +
                [f"{fmt_bytes(n)}：实得 GB/s" for n in sizes], rows)
    t_small = NVLINK4.time(4e3)
    print(f"\n例：在 NVLink 上发 4 KB 要 {fmt_time(t_small)}，其中带宽部分只有 {fmt_time(4e3 / NVLINK4.bw)}——"
          f"{NVLINK4.alpha / t_small * 100:.0f}% 的时间花在延迟上。")
    print_observations([
        f"链路越快，半带宽点越大：NVLink 要 {fmt_bytes(NVLINK4.half_bw_size())} 以上的消息才能用上一半带宽，"
        f"以太网只要 {fmt_bytes(ETH100.half_bw_size())}。带宽涨得比延迟降得快，'小消息'的门槛在不断提高。",
        "1 KB、64 KB 的消息在所有链路上都远远用不满带宽：它们是**延迟主导**的。LLM decode 时张量并行每层要 all-reduce 的"
        "激活只有几 KB 到几百 KB（05 脚本），正好落在这个区间——这时减少'步数'比减少'字节'更重要。",
        "1 GB 的消息在所有链路上都几乎跑满带宽（**带宽主导**）：prefill 的激活、KV cache 的搬运、权重加载属于这一类，"
        "这时'少传字节'才是关键。",
        "同一个 α–β 模型会贯穿本讲：每种集合通信算法的时间都可以写成 '几步 × α + 多少字节 / β'（03 脚本）。",
    ])


# ---------------------------------------------------------------------------
# 演示 2：拓扑
# ---------------------------------------------------------------------------
def demo2_topologies():
    """
    几种拓扑在 p 个设备时的性质。直径用模拟器的路由算出（p ≤ 64），对分带宽用公式（comm_sim.bisection_bw）。
    为了公平比较拓扑本身，所有拓扑的每条链路都取同样的带宽（ICI 的 45 GB/s）；两级网络用 NVLink + IB 的真实组合。
    """
    print_header("演示 2：拓扑——端口数、直径、对分带宽（同样的单链路带宽 45 GB/s，两级网络除外）")
    bw = ICI_V5E.bw
    rows = []
    for p in (8, 16, 64):
        a = int(round(p ** 0.5))
        torus_ok = a * a == p
        cases = [("交换机（无阻塞）", switch(p, ICI_V5E), "switch", 1, f"1 台 {p} 口交换机"),
                 ("双向环", ring(p, ICI_V5E), "ring", 2, "无")]
        if torus_ok:
            cases.append((f"{a}×{a} 二维 torus", torus2d(a, a, ICI_V5E), "torus", 4, "无"))
        for name, topo, kind, ports, sw in cases:
            bis = bisection_bw(kind, p, bw)
            rows.append([p, name, ports, topo.diameter(), f"{bis / 1e9:,.0f} GB/s", f"{bis / (p / 2) / 1e9:.1f} GB/s", sw])
    for n_nodes in (2, 8):
        p = 8 * n_nodes
        topo = two_level(n_nodes, 8)
        bis = bisection_bw("two_level", p, None, inter_bw=IB_NDR.bw)
        rows.append([p, f"两级：{n_nodes} 个节点 × 8 GPU", "1 NVLink + 1 网卡", topo.diameter(),
                     f"{bis / 1e9:,.0f} GB/s", f"{bis / (p / 2) / 1e9:.1f} GB/s", f"{n_nodes} 台节点交换机 + 胖树"])
    print_table(["设备数 p", "拓扑", "每设备端口", "直径（跳）", "对分带宽（单向）", "对分带宽 / 半数设备", "交换机"], rows)
    print_observations([
        "交换机：任意两个设备 2 跳可达（上行 + 下行），对分带宽随 p 线性增长——每个设备都能同时满速和另一半说话。"
        "代价是交换机本身：端口数随设备数增长，大交换机又贵又耗电。NVSwitch 把一个节点（8 块 H100，GB200 NVL72 是 72 块）连成这样。",
        "双向环：每个设备只有 2 个端口，便宜；但直径 p/2，对分带宽只有 2 条链路，不随 p 增长——设备一多，"
        "'人人对人人'的通信（all-to-all）就被中间那两条链路卡死（04 脚本）。",
        "二维 torus（TPU 的做法）：每个设备 4 个端口，没有交换机，规模可以做得很大；直径 √p，对分带宽随 √p 增长，"
        "介于环和交换机之间。好在 ring all-reduce 这类'只和邻居说话'的算法在 torus 上能跑满带宽（04 脚本）。",
        "两级网络（GPU 集群的做法）：节点内 NVLink 交换机很快，节点间每块 GPU 只有一块 50 GB/s 的网卡。"
        "跨节点的对分带宽按网卡算，每个 GPU 只有节点内的 1/9——所以张量并行这类频繁通信的切分方式一般限制在节点之内。",
    ])


# ---------------------------------------------------------------------------
# 演示 3：点对点
# ---------------------------------------------------------------------------
def demo3_point_to_point():
    """
    用模拟器发一条消息（设备 0 → 设备 1，或者跨节点的 0 → 8），看不同网络里的时间。
    """
    print_header("演示 3：点对点——从一块芯片发数据给另一块")
    cases = [("同一节点，NVLink 4", two_level(2, 8), 0, 1),
             ("跨节点，InfiniBand NDR", two_level(2, 8), 0, 8),
             ("PCIe 5.0（经过交换机）", switch(2, PCIE5), 0, 1),
             ("TPU v5e，相邻芯片（1 跳）", torus2d(16, 16, ICI_V5E), 0, 1),
             ("TPU v5e，对角最远（16 跳）", torus2d(16, 16, ICI_V5E), 0, 8 * 16 + 8),
             ("100 Gb 以太网", switch(2, ETH100), 0, 1)]
    rows = []
    for name, topo, s, d in cases:
        cells = [name, topo.hops(s, d)]
        for n in (4e3, 1e9):
            comm = Comm(topo)
            cells.append(fmt_time(comm.step([(s, d, n)])))
        rows.append(cells)
    print_table(["路径", "跳数", "4 KB", "1 GB"], rows)
    print_observations([
        "大消息（1 GB）只看路径上最慢的那条链路：跨节点比节点内慢 9 倍（50 vs 450 GB/s）；TPU 的跳数多，但大消息的时间几乎不变——"
        "数据像流水一样穿过中间的芯片（这里的模型假设是这样）。",
        "小消息（4 KB）几乎只看延迟：TPU 上走 16 跳的时间是 1 跳的十几倍，跳数开始要紧了。这就是为什么 torus 上要尽量只和邻居通信。",
        "注意这是一对一的发送；集合通信要协调很多设备，'几步'和'每步多少字节'才是关键（02、03 脚本）。",
    ])


if __name__ == "__main__":
    demo1_alpha_beta()
    demo2_topologies()
    demo3_point_to_point()