# 分块矩阵乘法

## 实现目标

`matmul.py` 使用 Triton 计算 `(M, K) @ (K, N) -> (M, N)`，重点展示分块矩阵
乘法、组主序 PID 映射和自动调优。输入会转换为 FP16，乘加过程使用 FP32
累加器，最终输出为 FP16。

## 核心数据流

1. `matmul()` 检查二维矩阵的形状兼容性，分配输出，并创建覆盖所有输出 tile 的
   一维启动网格。
2. `_matmul_kernel()` 将线性 PID 按组主序映射为 `(PID_M, PID_N)`，确定当前程序
   负责的输出 tile。
3. 内核生成 A、B 两个输入 tile 的地址，沿 K 维逐块加载并通过 `tl.dot()` 累加到
   FP32 寄存器。
4. 最后将累加器转换为 FP16，并使用二维边界掩码写入 C。
5. `triton.autotune` 根据 M、N、K 为当前 GPU 选择块大小、warp 数和流水线阶段数。

## 关键细节与限制

- 组主序让相邻 PID 更可能复用 A 的行块和 B 的列块，从而改善缓存局部性。
- K 维尾块使用掩码并以零填充；M、N 边界在最终存储时屏蔽。
- 包装函数仅支持二维、形状兼容的输入，并会将输入转换为 FP16。
- 自动调优首次遇到新的 M、N、K 组合时需要评估多组配置，因此会产生额外开销。

## 测试与基准

```bash
python3 matmul.py
python3 matmul.py --benchmark
```

数值测试以 `torch.matmul()` 为参考；基准测试比较两种实现的 TFLOPS。

## 配套视频

[bilibili](https://www.bilibili.com/video/BV1fMyWBgERM?spm_id_from=333.788.videopod.episodes&vd_source=e98b669ccbafff4b5aa59dd6303b722f&p=9)
