By name: pypto-api description: "PyPTO 全部 API 签名与约束速查" category: fundamental version: "1.0.0" metadata: backend: ascend dsl: pypto
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):
...
张量
| API | 用途 | 示例 |
|---|---|---|
pypto.Tensor(shape, dtype) |
输入/输出类型标注 | x: pypto.Tensor((m, k), pypto.DT_FP32) |
pypto.tensor(shape_list, dtype) |
kernel 内创建输出 | output = pypto.tensor([m, n], pypto.DT_FP32) |
pypto.zeros(shape_list, dtype=) |
零初始化张量（累加器） | acc = pypto.zeros([1], dtype=pypto.DT_FP32) |
pypto.full(shape, val, dtype, valid_shape=) |
常量填充张量 | ones = pypto.full(s, 1.0, pypto.DT_FP32, valid_shape=s) |
数据类型：pypto.DT_FP32、pypto.DT_INT32、pypto.DT_INT64（INT64 仅用于输入标注）。
Tile 配置
| API | 约束 |
|---|---|
pypto.set_vec_tile_shapes(*shapes) |
参数个数 = 被操作张量 rank |
pypto.set_cube_tile_shapes(m, k, n, l1, split_k) |
固定 5 参数 |
tile 双约束：
prod(tile_shape)≤ 16384auto_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)
分块
| API | 用途 | 约束 |
|---|---|---|
pypto.loop(start, end, step, name=, idx_name=) |
编译期循环 | 不嵌套、尽量少用 |
pypto.view(tensor, shape, offset) |
切片提取（等价 tensor[a:b, c:d]） |
shape 各维为编译期常量，rank 不变，每维 ≤ 输入对应维 |
pypto.assemble(chunk, offset, output) |
写回子块（等价切片赋值） | 无 |
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
数学函数
| 函数 | 说明 |
|---|---|
pypto.exp(x) |
指数 |
pypto.log(x) |
对数 |
pypto.sqrt(x) |
平方根 |
pypto.abs(x) |
绝对值 |
pypto.sigmoid(x) |
sigmoid |
pypto.softmax(x, dim=) |
softmax |
pypto.maximum(a, b) |
逐元素最大，b 可以是标量：pypto.maximum(x, 0.0) |
pypto.minimum(a, b) |
逐元素最小，b 可以是标量：pypto.minimum(x, 0.0) |
有内建函数时直接调用，禁止手写等价公式。
禁用 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。
类型转换与索引
| API | 示例 |
|---|---|
pypto.cast(tensor, dtype) |
pypto.cast(x, pypto.DT_INT32) |
pypto.unsqueeze(tensor, dim) |
pypto.unsqueeze(x, 1) |
pypto.gather(tensor, dim, index) |
pypto.gather(log_probs, dim=1, index=idx) |
pypto.expand_clone(tensor, shape) |
单轴广播，每次只能扩展一个轴 |