pypto-basics

PyPTO 编程原则

原则 0：先抽取 Torch 基线的语义合同

写 PyPTO 前，先把 baseline forward 变成“可执行语义合同”，再做实现与优化。

最少做 5 件事：

1. 写出数学式（谁参与计算、谁是被规约对象、输出定义是什么）。
2. 核对 API 语义而不是变量名（input/target/prediction 这些名字经常误导）。
3. 固化规约合同（sum、batchmean、或 mean 语义=sum/count，并明确规约轴与 keepdim）。
4. 判断是否非对称（A||B 与 B||A 是否等价）。
5. 做 1 组语义自检样例（优先用非对称输入，先验方向是否正确）。

注意：性能规则（tile/loop）只能在语义合同确定后应用，不能反向决定语义。

原则 1：静态 shape

kernel 一切在编译时确定。工厂函数封装 kernel，shape 和标量参数（eps、slope 等）作为闭包传入。forward 直接传 torch.Tensor。ModelNew.init 签名必须与原始 Model 一致，shape 在 forward 中获取。

补充（高优先级）：

• benchmark 的 get_inputs/get_init_inputs 在单次任务里是固定参数；把它当静态合同，不要过度通用化。
• 先做“固定参数抽取”：从题目文件读取本次 get_init_inputs 的返回值，作为本次任务常量（例如 dim=1）。
• 题目里的注释如 Example, change to desired ... 视为数据集说明噪声，不是本次实现要求。
• 对固定 dim 的任务，生成单一固定 dim kernel；不要在一个 kernel 里写 if dim == 0/1/2 分支。
• 若确实要支持多个 dim，使用多个工厂函数/多个 kernel 分别生成，不要把多语义揉进同一 kernel。

原则 2：forward 的职责

1. Assert：assert x.dim() == N、assert tuple(x.shape) == (...)
2. Reshape（如需）：torch reshape 为 kernel 能处理的 shape
3. 调用 kernel：传入 contiguous 张量，reshape 回原始 shape 返回

forward 内禁止 torch 计算。kernel 内不能 reshape。 输出语义（keepdim/是否 squeeze）要与 baseline 直接对齐；不要先改语义再靠额外 squeeze/unsqueeze 回补。

原则 3：选择维度

算子类型	forward 策略	kernel 维度
Elementwise / 简单 Loss	reshape(-1)	1D
GroupNorm / InstanceNorm	reshape(flat_batch, hidden)	2D
BatchNorm / RMSNorm	reshape(B, C, -1)	3D
Batched matmul（同 batch）	保持 3D，不需要 loop	3D
2D Matmul	保持 2D	2D
单轴归约	保持原始维度	原维

补充：

• Elementwise 算子无数据依赖，优先 flatten 为 1D；loop 仅由 auto_tiles > 2048 触发，与维度无关。
• 保持高维的唯一理由是业务语义需要（如后续算子依赖布局），而非 tile 约束。

补充：

• 多轴规约若轴连续且中间结果不被其他算子使用，优先在 forward 合并为单轴后再规约（例如 H,W -> HW）。

原则 4：tile 双约束

1. prod(tile_shape) ≤ 16384
2. auto_tiles = prod(ceil(shape[i]/tile[i])) ≤ 2048

tile 参数个数 = 被操作 tensor 的 rank。常用：

• 1D: (8192) | 2D: (1, 16384) | 3D: (1, 1, 16384) 或 (1, 16, 256)（单轴 3D 归约常见起步）
• matmul: set_cube_tile_shapes([128, 128], [32, 128], [256, 256], True, False)
• 核心红线：Skill/示例中的 shape、tile、BLOCK 常量（如 16384/8192/4096）只能作参考。必须按当前任务的输入维度重新计算，禁止直接照抄。
• 经验规则：优先避免明显 tile[i] > shape[i] 的参数浪费；如确需使用，必须有明确理由（例如实测收益）。