PyPTO Pass
PyPTO Pass 遵循了 MLIR 的多层 IR 渐进式 lowering 设计思路,将 PyPTO 前端图到 NPU 码的编译路径分为三层,逐层降低抽象、增加硬件感知。与 vLLM 等推理框架对编译栈的分层设计目标一致——在合适层级暴露优化机会,同时隔离复杂度:
Tensor Graph 层
完成 Python 侧算子 API 到计算图 IR 的展开与落地,将异构 dtype 统一、冗余视图消除,产出规范化的算子级图表示,确保后续 Pass 输入一致可预测。
Tile Graph 层
引入 NPU 内存层级和分块语义,对连续视图链合并折叠、对计算任务并行拆分,通过显式标注数据在 DDR/L1/L0/UB 间搬移路径,将平台无关算子转化为硬件感知的分块图,兼顾访存效率与并行吞吐。
Block Graph 层
以实现核内核间计算优化为目的,完成子图函数化、全局内存复用、核内乱序调度与同步,补齐动态 Shape 和运行时属性,产出可被 Codegen 直接消费的最终 IR。
算子代码
def add_mul_kernel(
a: pypto.Tensor([], pypto.DT_FP32),
b: pypto.Tensor([], pypto.DT_FP32),
d: pypto.Tensor([], pypto.DT_FP32),
c: pypto.Tensor([], pypto.DT_FP32),
e: pypto.Tensor([], pypto.DT_FP32),
):
# Vector tile 16x16 covers element-wise add and mul on 32x32 tensors
pypto.set_vec_tile_shapes(16, 16)
# Cube tile 16x16 (M=16, K=16, N=16) — declared per spec
pypto.set_cube_tile_shapes([16, 16], [16, 16], [16, 16])
c[:] = pypto.add(a, b) # c = a + b
e[:] = pypto.mul(d, c) # e = d * c
return add_mul_kernel
张量计算图
Loading computation graph...
Tile切分
Loading computation graph...
优化与染色
Loading computation graph...