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...