Pypto API

PyPTO API 速查

Kernel 装饰器

@pypto.frontend.jit(
    runtime_options={"run_mode": _PYPTO_RUN_MODE},
    debug_options={"runtime_debug_mode": _PYPTO_RUNTIME_DEBUG_MODE},
)
def kernel(x: pypto.Tensor(shape_tuple, dtype)) -> pypto.Tensor(shape_tuple, dtype):
    ...

张量

数据类型：pypto.DT_FP32、pypto.DT_INT32、pypto.DT_INT64（INT64 仅用于输入标注）。

Tile 配置

tile 双约束：

1. prod(tile_shape) ≤ 16384
2. auto_tiles = prod(每维 ceil(shape[i]/tile[i])) ≤ 2048（每个 op）

• 必须在任何计算操作之前调用。一个 kernel 内可多次切换 tile。
• 若 auto_tiles > 2048，优先改为 loop + view/assemble 分块实现。
• vec tile 推荐：(8192) (1D)、(1, 16384) (2D)、(1, 1, 16384) (3D)
• cube tile 推荐：set_cube_tile_shapes([128, 128], [32, 128], [256, 256], True, False)

分块

pypto.view 不是 reshape。它是 tensor[offset[0]:offset[0]+shape[0], ...] 的等价 API。不能改变维度数，不能改变维度排布。所有 reshape 必须在 forward 中用 torch 完成。

算术运算

运算符规则：+ * 支持标量在任意位置；- / 要求 tensor 在左侧（1.0 - x crash）。 函数调用：pypto.add/sub/mul/div 第一参数必须 Tensor。 一元取反 -x：不支持，用 pypto.mul(x, -1.0) 或 x * (-1.0)。 切片赋值：output[:] = expr

数学函数

有内建函数时直接调用，禁止手写等价公式。

禁用 API：pypto.where（有 bug）、pypto.clamp（不稳定）。条件逻辑用 maximum/minimum 组合实现。

归约

pypto.sum(x, dim=int, keepdim=bool)
pypto.amax(x, dim=int, keepdim=bool)
pypto.amin(x, dim=int, keepdim=bool)

无 pypto.mean API。 mean 语义请用 sum * (1.0 / count) 实现。 dim 只接受单个 int，不支持 list。 多轴归约需连续多次调用。 dim 在实际实现中应作为闭包常量参与编译；静态任务不要在同一 kernel 内写多 dim 运行时分支。

矩阵乘法

pypto.matmul(a, b, out_dtype, a_trans=False, b_trans=False)

• 2D：[M,K] @ [K,N] → [M,N]
• 3D batched：[B,M,K] @ [B,K,N] → [B,M,N]（两边必须同 rank 同 batch，不支持广播）
• a_trans / b_trans 支持转置。
• 限制：每个输入最后一维 ≤ 65535。超过时用 sum(a * b_broadcast, dim=) 替代。
• matmul 几乎必须配合 loop（M 轴分块），因 matmul 内部展开复杂，不 loop 容易 timeout。例外：M 很小（≤128）时可不 loop。

类型转换与索引